153 lines
6.2 KiB
Markdown
153 lines
6.2 KiB
Markdown
# laoyao_2b_moe
|
||
|
||
这个 repo 用来把 Jiayi/Laoyao 的 2B MoE 小模型实验重构成更可复现的预训练项目：数据配比更合理，模型定义迁移到 NeMo/Megatron 风格配置，训练入口从手写 PyTorch 逻辑迁移到 Megatron/Nemo 体系。
|
||
|
||
当前目标不是继续沿用旧的手写 trainer，而是把数据、模型、训练、评测和导出几个边界拆开，形成一个可以反复运行和审计的 ML project。
|
||
|
||
## 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`。
|
||
|
||
## 当前数据计划
|
||
|
||
预训练目标规模为 200B tokens，随机种子固定为 42。核心配比：
|
||
|
||
| 类别 | 目标比例 | 说明 |
|
||
|---|---:|---|
|
||
| English Web | 40% | Ultra-FineWeb/FineWeb 类英文网页文本 |
|
||
| English Edu | 20% | FineWeb-Edu/高质量教育文本；science 不足时也从这里补齐 |
|
||
| Chinese Clean | 10% | Ultra-FineWeb zh / CCI3-HQ / SkyPile 类中文干净文本 |
|
||
| Science | 10% | medmcqa、proofwriter、scienceqa、sciq、qasc、openbookqa；不足则 edu 补齐 |
|
||
| Logic | 10% | 从高质量 scored/education 数据中抽取逻辑类文本 |
|
||
| Math | 5% | 使用 Jiayi 已清洗/评分的数学来源 |
|
||
| Code | 5% | 使用 Jiayi 已去重的多语言代码数据 |
|
||
|
||
验证集为 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 设计
|
||
|
||
为了兼容 GLM-5.2 生态，本模型的 tokenizer/vocab 直接采用 GLM-5.2 同款词表，而不是旧 Laoyao 65K BPE。当前 g0033 上的引用路径是：
|
||
|
||
```bash
|
||
/mnt/beegfs/yi/laoyao_2b_moe/tokenizer/glm5.2/tokenizer.json
|
||
```
|
||
|
||
该 tokenizer 从 ModelScope zai-org/GLM-5.2 下载，当前训练配置使用的 vocab size 为 154,856，tokenizer.json sha256 为 19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d。后续训练数据应按这个 tokenizer 重新统计/切分 token budget。
|
||
|
||
## g0050 训练状态入口
|
||
|
||
g0050 上的实际训练 repo 路径：
|
||
|
||
```bash
|
||
/ssd/workspace/yi/laoyao_2b_moe
|
||
```
|
||
|
||
当前主训练入口为：
|
||
|
||
```bash
|
||
bash scripts/resume_pretrain_8192_8gpu_mbs14.sh
|
||
```
|
||
|
||
该入口默认从 `runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints` 恢复，使用：
|
||
|
||
- `seq_length=8192`
|
||
- `micro_batch_size=14`
|
||
- `global_batch_size=112`
|
||
- `tensor/pipeline/expert/context parallel = 1/1/1/1`
|
||
- `distributed optimizer + overlap grad reduce + overlap param gather`
|
||
- `full recompute, uniform, recompute_num_layers=1`
|
||
- `save_interval=2500`
|
||
- `eval_interval=15000`
|
||
- heldout validation prefix: `dataset/val/megatron_8192_glm52/heldout_2p8k_text_document`
|
||
|
||
Megatron indexed validation 数据是本地派生产物，被 `.gitignore` 忽略；需要在目标机器上从 `dataset/val/data/heldout_2p8k_packed_text.jsonl` 重新生成或从训练机器同步。
|
||
|
||
## g0033 路径约定
|
||
|
||
repo 路径：
|
||
|
||
```bash
|
||
/mnt/beegfs/yi/laoyao_2b_moe
|
||
```
|
||
|
||
预训练数据实际读取路径：
|
||
|
||
```bash
|
||
/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701
|
||
```
|
||
|
||
数据不再复制到 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` 作为辅助指标，因为它更容易偏向短选项。
|