diff --git a/AGENTS.md b/AGENTS.md index ed45450..a7ef219 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -289,6 +289,40 @@ runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000 Do not commit TensorBoard logs. +## HF Heldout 2.8k Evaluation + +After exporting a checkpoint to HF format, run: + +```bash +python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --data dataset/val/data/heldout_2p8k_sft_prompt_completion.jsonl \ + --out-dir runs/hf_eval/heldout_2p8k/iter_0107500 \ + --device cuda \ + --max-length 2048 +``` + +For CPU smoke tests, use a very small sample count: + +```bash +python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --max-items 4 \ + --device cpu \ + --dtype float32 \ + --max-length 512 +``` + +This evaluation follows the tokenizer-swap heldout logic: + +- Compute prompt-conditioned completion NLL/PPL for all 2.8k examples. +- Parse MCQ choices only where the prompt has explicit `A.`/`B.`/... choices and the gold completion starts with the answer label. +- Score each MCQ choice as a continuation. +- Main accuracy is `mcq_acc_avg_norm`, based on average token logprob. +- `mcq_acc_sum` is auxiliary and can favor shorter choices. + +Do not treat all 2.8k examples as MCQ. The set also includes GSM8K, HumanEval/MBPP style code, Chinese dialogue, and English dialogue examples. + ## Model Weight Exporting The HF export tools are under: diff --git a/README.md b/README.md index c4cbb64..2cb3e4c 100644 --- a/README.md +++ b/README.md @@ -2,14 +2,23 @@ 这个 repo 用来把 Jiayi/Laoyao 的 2B MoE 小模型实验重构成更可复现的预训练项目：数据配比更合理，模型定义迁移到 NeMo/Megatron 风格配置，训练入口从手写 PyTorch 逻辑迁移到 Megatron/Nemo 体系。 -当前目标不是继续沿用旧的手写 trainer，而是把数据、模型、训练三个边界拆开： +当前目标不是继续沿用旧的手写 trainer，而是把数据、模型、训练、评测和导出几个边界拆开，形成一个可以反复运行和审计的 ML project。 -- `dataset/`: 预训练数据和验证集的构建、manifest、数据落盘位置。 -- `model/`: 2B MoE 架构定义，优先用 NeMo/Megatron 配置表达。 -- `training/`: 训练 recipe、评估 recipe、并行和优化超参。 -- `scripts/`: 数据下载、tokenization、训练恢复、Megatron 推理服务和 HF 导出 smoke 的 shell 入口。 -- `tools/`: checkpoint 检查、Megatron DCP 探测、Megatron-to-HF 导出和 HF custom model 推理工具。 -- `docs/`: 训练/导出过程中沉淀的问题记录。 +## Repo 结构 + +| 路径 | 功能 | +|---|---| +| `dataset/pretrain/` | 200B 级预训练数据的下载、清洗、配比、manifest 和 parquet -> Megatron indexed dataset 转换。 | +| `dataset/val/` | heldout 2.8k 验证集源数据和 Megatron validation 数据说明。 | +| `model/` | 2B MoE 架构定义和 Megatron/NeMo 配置说明。 | +| `training/` | Megatron-Bridge 训练 recipe、模型构建参数、数据 blend、validation、并行和优化配置。 | +| `docker/` | NeMo/Megatron 训练镜像，当前使用 flash-attn4 修复后的镜像。 | +| `scripts/` | 数据下载、tokenization、镜像构建、训练启动/恢复、Megatron inference server、查询 smoke。 | +| `tools/` | checkpoint 探测、DCP 元数据检查、Megatron -> HF 导出、HF custom model 生成和 heldout eval。 | +| `docs/` | 训练、推理、导出过程中遇到的问题和修复记录。 | +| `runs/` | checkpoint、TensorBoard、HF export、eval 输出等运行产物；默认不进 git。 | + +给 agent 的更详细操作手册见根目录 `AGENTS.md`。 ## 当前数据计划 @@ -27,6 +36,35 @@ 验证集为 2,800 条，七个能力维度各 400 条：science_reasoning、logic、code、chinese_exam、math、chinese_dialogue、english_dialogue。它会用于预训练过程中的 accuracy、perplexity 和 nll 评估。 +## 当前主训练配置 + +g0050 上当前主训练入口： + +```bash +cd /ssd/workspace/yi/laoyao_2b_moe +bash scripts/resume_pretrain_8192_8gpu_mbs14.sh +``` + +核心配置： + +- Megatron-Bridge / Megatron-Core backend。 +- Docker image: `laoyao/nemo-megatron:26.06-flashattn4`。 +- `seq_length=8192`。 +- `micro_batch_size=14`。 +- `global_batch_size=112`。 +- `DP=8, TP=1, PP=1, EP=1, CP=1`。 +- distributed optimizer 开启。 +- grad reduce / param gather overlap 开启。 +- full activation recompute: `uniform`, `recompute_num_layers=1`。 +- 每 `2500` iter 保存 checkpoint，保留最近 `10` 个。 +- 每 `15000` iter 做 heldout validation，`eval_iters=10`。 + +当前 run 名称： + +```text +pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000 +``` + ## Tokenizer 设计 @@ -81,3 +119,34 @@ repo 路径： ``` 数据不再复制到 repo 内；训练直接读取上面的源输出目录。`dataset/pretrain/data/` 仍被 `.gitignore` 忽略，仅用于小规模临时样本。 + +## 模型导出与 heldout 评测 + +Megatron DCP checkpoint 可导出为 HuggingFace custom model： + +```bash +python3 tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py \ + --checkpoint-dir runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints/iter_0107500 \ + --tokenizer-dir tokenizer/glm5.2 \ + --output-dir runs/hf_exports/iter_0107500 +``` + +HF custom model 当前没有实现 KV cache，因此生成和评测脚本均显式 `use_cache=False`。 + +heldout 2.8k HF 评测入口： + +```bash +python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --data dataset/val/data/heldout_2p8k_sft_prompt_completion.jsonl \ + --out-dir runs/hf_eval/heldout_2p8k/iter_0107500 \ + --device cuda \ + --max-length 2048 +``` + +评测逻辑沿用 tokenizer-swap 实验的口径： + +- 所有样本计算 prompt-conditioned completion NLL/PPL、NLL/token、bits/byte。 +- 对能解析出候选项的 MCQ 样本，按每个候选 continuation 的 logprob 打分。 +- 主 accuracy 是平均 token logprob 归一化后的 `mcq_acc_avg_norm`。 +- `mcq_acc_sum` 作为辅助指标，因为它更容易偏向短选项。 diff --git a/tools/hf_laoyao_moe/README.md b/tools/hf_laoyao_moe/README.md index f133a22..86ab838 100644 --- a/tools/hf_laoyao_moe/README.md +++ b/tools/hf_laoyao_moe/README.md @@ -32,6 +32,37 @@ python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \ 当前 HF custom model 没有实现 KV cache，因此生成脚本强制 `use_cache=False`。不要删除这个设置；否则 Transformers 默认 cache path 会导致上下文丢失和异常重复。 +## Heldout 2.8k 评测 + +`eval_heldout_2p8k.py` 复用 tokenizer-swap 实验的评测口径： + +- 对所有样本计算 prompt-conditioned completion NLL/PPL。 +- 对能从 prompt 中解析出 A/B/C/D 候选项、且 gold completion 以候选 label 开头的样本计算 MCQ accuracy。 +- 主 accuracy 是候选 continuation 的平均 token logprob accuracy：`mcq_acc_avg_norm`。 +- `mcq_acc_sum` 是总 logprob accuracy，容易偏好短答案，只作辅助指标。 + +示例： + +```bash +python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --data dataset/val/data/heldout_2p8k_sft_prompt_completion.jsonl \ + --out-dir runs/hf_eval/heldout_2p8k/iter_0107500 \ + --device cuda \ + --max-length 2048 +``` + +CPU smoke 可以限制样本数： + +```bash +python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \ + --model-dir runs/hf_exports/iter_0107500 \ + --max-items 4 \ + --device cpu \ + --dtype float32 \ + --max-length 512 +``` + ## 调试 打印逐 token 结果： @@ -66,4 +97,3 @@ python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \ - tied embedding/output head - MoE top-k router with `moe_expert_capacity_factor=1.25` - `use_cache=false` - diff --git a/tools/hf_laoyao_moe/eval_heldout_2p8k.py b/tools/hf_laoyao_moe/eval_heldout_2p8k.py new file mode 100755 index 0000000..e24d2df --- /dev/null +++ b/tools/hf_laoyao_moe/eval_heldout_2p8k.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +import math +import re +from collections import defaultdict +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F +from transformers import AutoModelForCausalLM, AutoTokenizer + + +CHOICE_RE = re.compile(r"(?m)^\s*([A-Z])\s*[\.|\)]\s+(.+?)(?=\n\s*[A-Z]\s*[\.|\)]\s+|\Z)", re.DOTALL) +ANSWER_LABEL_RE = re.compile(r"^\s*([A-Z])\s*[\.|\)]") + + +def load_jsonl(path: Path, max_items: int = 0, num_shards: int = 1, shard_id: int = 0) -> list[dict[str, Any]]: + rows = [] + with path.open("r", encoding="utf-8") as handle: + for idx, line in enumerate(handle): + if not line.strip(): + continue + if num_shards > 1 and idx % num_shards != shard_id: + continue + rows.append(json.loads(line)) + if max_items and len(rows) >= max_items: + break + return rows + + +def encode(tokenizer, text: str) -> list[int]: + return tokenizer.encode(text, add_special_tokens=False) + + +@torch.inference_mode() +def continuation_nll(model, tokenizer, prompt: str, continuation: str, max_length: int) -> dict[str, float] | None: + prompt_ids = encode(tokenizer, prompt) + cont_ids = encode(tokenizer, continuation) + if not cont_ids: + return None + + if len(cont_ids) >= max_length: + cont_ids = cont_ids[: max_length - 1] + prompt_ids = [] + else: + prompt_budget = max_length - len(cont_ids) + prompt_ids = prompt_ids[-prompt_budget:] + + ids = prompt_ids + cont_ids + labels = [-100] * len(prompt_ids) + cont_ids + if len(ids) < 2: + return None + + input_ids = torch.tensor([ids], device=model.device, dtype=torch.long) + label_ids = torch.tensor([labels], device=model.device, dtype=torch.long) + logits = model(input_ids=input_ids, use_cache=False).logits + + shift_logits = logits[:, :-1, :].contiguous() + shift_labels = label_ids[:, 1:].contiguous() + mask = shift_labels.ne(-100) + if not mask.any(): + return None + + log_probs = F.log_softmax(shift_logits.float(), dim=-1) + safe_labels = shift_labels.masked_fill(~mask, 0) + token_log_probs = log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1) + nll = -token_log_probs[mask].sum().item() + tokens = int(mask.sum().item()) + return {"nll": nll, "tokens": tokens} + + +def score_completion(model, tokenizer, prompt: str, completion: str, max_length: int) -> dict[str, float] | None: + ret = continuation_nll(model, tokenizer, prompt, completion, max_length) + if ret is None: + return None + byte_len = max(1, len(completion.encode("utf-8"))) + ret["nll_per_token"] = ret["nll"] / max(1, ret["tokens"]) + ret["ppl"] = math.exp(min(50.0, ret["nll_per_token"])) + ret["nll_per_byte"] = ret["nll"] / byte_len + ret["bits_per_byte"] = ret["nll"] / (byte_len * math.log(2)) + ret["bytes"] = byte_len + return ret + + +def parse_mcq(prompt: str, completion: str) -> dict[str, Any] | None: + matches = list(CHOICE_RE.finditer(prompt)) + if len(matches) < 2: + return None + + labels = [m.group(1) for m in matches] + choices = [f"{m.group(1)}. {m.group(2).strip()}" for m in matches] + answer_match = ANSWER_LABEL_RE.match(completion) + if answer_match is None: + return None + answer_label = answer_match.group(1) + if answer_label not in labels: + return None + + first_choice_start = matches[0].start() + mcq_prompt = prompt[:first_choice_start].rstrip() + if not mcq_prompt: + mcq_prompt = prompt.rstrip() + return { + "mcq_prompt": mcq_prompt, + "labels": labels, + "choices": choices, + "answer_idx": labels.index(answer_label), + } + + +def score_mcq(model, tokenizer, prompt: str, choices: list[str], max_length: int) -> tuple[list[dict[str, float]], int, int]: + scores = [] + for choice in choices: + sep = "" if prompt.endswith((" ", "\n")) else "\n" + ret = continuation_nll(model, tokenizer, prompt + sep, choice, max_length) + if ret is None: + scores.append({"sum_logprob": -float("inf"), "avg_logprob": -float("inf"), "tokens": 0}) + continue + sum_logprob = -ret["nll"] + avg_logprob = sum_logprob / max(1, ret["tokens"]) + scores.append({"sum_logprob": sum_logprob, "avg_logprob": avg_logprob, "tokens": ret["tokens"]}) + pred_avg = max(range(len(scores)), key=lambda idx: scores[idx]["avg_logprob"]) + pred_sum = max(range(len(scores)), key=lambda idx: scores[idx]["sum_logprob"]) + return scores, pred_avg, pred_sum + + +def mean(values: list[float | int | None]) -> float | None: + xs = [float(x) for x in values if x is not None] + return sum(xs) / len(xs) if xs else None + + +def summarize(results: list[dict[str, Any]]) -> dict[str, Any]: + groups: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in results: + groups[row["capability"]].append(row) + groups["all"] = results + + out = {} + for name, rows in sorted(groups.items()): + mcq_rows = [row for row in rows if row.get("mcq") is not None] + ppl_rows = [row for row in rows if row.get("ppl") is not None] + out[name] = { + "n": len(rows), + "mcq_n": len(mcq_rows), + "mcq_acc_avg_norm": mean([row["mcq"]["correct_avg_norm"] for row in mcq_rows]), + "mcq_acc_sum": mean([row["mcq"]["correct_sum"] for row in mcq_rows]), + "ppl_n": len(ppl_rows), + "ppl_token_mean": mean([row["ppl"]["ppl"] for row in ppl_rows]), + "nll_per_token_mean": mean([row["ppl"]["nll_per_token"] for row in ppl_rows]), + "bits_per_byte_mean": mean([row["ppl"]["bits_per_byte"] for row in ppl_rows]), + "completion_tokens_mean": mean([row["ppl"]["tokens"] for row in ppl_rows]), + } + return out + + +def main() -> None: + parser = argparse.ArgumentParser(description="Evaluate Laoyao HF export on the heldout 2.8k validation set.") + parser.add_argument("--model-dir", required=True, help="HF export directory produced by convert_laoyao_dcp_to_hf.py") + parser.add_argument("--data", default="dataset/val/data/heldout_2p8k_sft_prompt_completion.jsonl") + parser.add_argument("--out-dir", default="runs/hf_eval/heldout_2p8k") + parser.add_argument("--model-label", default=None) + parser.add_argument("--max-items", type=int, default=0) + parser.add_argument("--max-length", type=int, default=2048) + parser.add_argument("--num-shards", type=int, default=1) + parser.add_argument("--shard-id", type=int, default=0) + parser.add_argument("--device", default="cuda") + parser.add_argument("--dtype", choices=["auto", "float16", "bfloat16", "float32"], default="bfloat16") + parser.add_argument("--fix-mistral-regex", action="store_true") + args = parser.parse_args() + + dtype = { + "auto": "auto", + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, + }[args.dtype] + + tokenizer = AutoTokenizer.from_pretrained( + args.model_dir, + trust_remote_code=True, + fix_mistral_regex=args.fix_mistral_regex, + ) + model = AutoModelForCausalLM.from_pretrained( + args.model_dir, + trust_remote_code=True, + torch_dtype=dtype, + ) + model.to(args.device) + model.eval() + if hasattr(model.config, "use_cache"): + model.config.use_cache = False + + rows = load_jsonl(Path(args.data), args.max_items, args.num_shards, args.shard_id) + label = args.model_label or Path(args.model_dir).name + out_label = label + if args.num_shards > 1: + if args.shard_id < 0 or args.shard_id >= args.num_shards: + raise ValueError("--shard-id must be in [0, --num-shards)") + out_label = f"{label}.shard{args.shard_id:02d}of{args.num_shards:02d}" + + results = [] + for idx, row in enumerate(rows, 1): + prompt = row["prompt"] + completion = row["completion"] + ppl = score_completion(model, tokenizer, prompt, completion, args.max_length) + + mcq_result = None + mcq = parse_mcq(prompt, completion) + if mcq is not None: + scores, pred_avg, pred_sum = score_mcq(model, tokenizer, mcq["mcq_prompt"], mcq["choices"], args.max_length) + mcq_result = { + "answer_idx": mcq["answer_idx"], + "labels": mcq["labels"], + "pred_avg_norm": pred_avg, + "pred_sum": pred_sum, + "correct_avg_norm": int(pred_avg == mcq["answer_idx"]), + "correct_sum": int(pred_sum == mcq["answer_idx"]), + "scores": scores, + } + + results.append( + { + "id": row.get("hashes", {}).get("pair_sha256") or f"row_{idx}", + "capability": row.get("capability", "unknown"), + "source_id": row.get("source_id", ""), + "split": row.get("split", ""), + "ppl": ppl, + "mcq": mcq_result, + } + ) + if idx % 100 == 0: + print(f"[{label}] evaluated {idx}/{len(rows)}", flush=True) + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + summary = { + "model_dir": args.model_dir, + "model_label": label, + "data": args.data, + "max_items": args.max_items, + "max_length": args.max_length, + "num_shards": args.num_shards, + "shard_id": args.shard_id, + "summary": summarize(results), + } + with (out_dir / f"{out_label}.per_item.jsonl").open("w", encoding="utf-8") as handle: + for item in results: + handle.write(json.dumps(item, ensure_ascii=False) + "\n") + with (out_dir / f"{out_label}.summary.json").open("w", encoding="utf-8") as handle: + json.dump(summary, handle, ensure_ascii=False, indent=2) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main()