Add heldout HF evaluation workflow

This commit is contained in:
yi_lu
2026-07-09 08:39:02 +08:00
parent e1e880cadf
commit 07e108388d
4 changed files with 400 additions and 8 deletions

View File

@@ -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` 作为辅助指标,因为它更容易偏向短选项。