Add heldout HF evaluation workflow

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:

README.md

@@ -2,14 +2,23 @@
 
 这个 repo 用来把 Jiayi/Laoyao 的 2B MoE 小模型实验重构成更可复现的预训练项目：数据配比更合理，模型定义迁移到 NeMo/Megatron 风格配置，训练入口从手写 PyTorch 逻辑迁移到 Megatron/Nemo 体系。
 
-当前目标不是继续沿用旧的手写 trainer，而是把数据、模型、训练三个边界拆开：
+当前目标不是继续沿用旧的手写 trainer，而是把数据、模型、训练、评测和导出几个边界拆开，形成一个可以反复运行和审计的 ML project。
 
-- `dataset/`: 预训练数据和验证集的构建、manifest、数据落盘位置。
+## Repo 结构
-- `model/`: 2B MoE 架构定义，优先用 NeMo/Megatron 配置表达。
+
-- `training/`: 训练 recipe、评估 recipe、并行和优化超参。
+| 路径 | 功能 |
-- `scripts/`: 数据下载、tokenization、训练恢复、Megatron 推理服务和 HF 导出 smoke 的 shell 入口。
+|---|---|
-- `tools/`: checkpoint 检查、Megatron DCP 探测、Megatron-to-HF 导出和 HF custom model 推理工具。
+| `dataset/pretrain/` | 200B 级预训练数据的下载、清洗、配比、manifest 和 parquet -> Megatron indexed dataset 转换。 |
-- `docs/`: 训练/导出过程中沉淀的问题记录。
+| `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` 作为辅助指标，因为它更容易偏向短选项。

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`
-

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()