laoyao_2b_moe

这个 repo 用来把 Jiayi/Laoyao 的 2B MoE 小模型实验重构成更可复现的预训练项目:数据配比更合理,模型定义迁移到 NeMo/Megatron 风格配置,训练入口从手写 PyTorch 逻辑迁移到 Megatron/Nemo 体系。

当前目标不是继续沿用旧的手写 trainer而是把数据、模型、训练三个边界拆开

  • 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/: 训练/导出过程中沉淀的问题记录。

当前数据计划

预训练目标规模为 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 评估。

Tokenizer 设计

为了兼容 GLM-5.2 生态,本模型的 tokenizer/vocab 直接采用 GLM-5.2 同款词表,而不是旧 Laoyao 65K BPE。当前 g0033 上的引用路径是:

/mnt/beegfs/yi/laoyao_2b_moe/tokenizer/glm5.2/tokenizer.json

该 tokenizer 从 ModelScope zai-org/GLM-5.2 下载,当前训练配置使用的 vocab size 为 154,856tokenizer.json sha256 为 19e773648cb4e65de8660ea6365e10acca112d42a854923df93db4a6f333a82d。后续训练数据应按这个 tokenizer 重新统计/切分 token budget。

g0050 训练状态入口

g0050 上的实际训练 repo 路径:

/ssd/workspace/yi/laoyao_2b_moe

当前主训练入口为:

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 路径:

/mnt/beegfs/yi/laoyao_2b_moe

预训练数据实际读取路径:

/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 忽略,仅用于小规模临时样本。

Description
The repository hosting jiayi (laoyao)'s 2B-MoE model definition, corresponding training data and relevant scripts
Readme 8 MiB
Languages
Python 75.1%
Shell 24.9%