laoyao_2b_moe/training/megatron_bridge/README.md

Megatron-Bridge Training

本目录是 Laoyao 2B MoE 在 NVIDIA Megatron-Bridge 上的训练适配层。

当前策略：

• 不改变原模型参数规模：hidden_size=1536、12 experts、topk=4、5 个 MoE layer。
• 训练上下文先用 seq_length=8192，不要一开始上 16K。
• tokenizer 使用 repo 内已验证的 GLM-5.2 tokenizer：tokenizer/glm5.2。
• 从零预训练使用 Megatron indexed dataset，不能直接把 parquet 喂给 Bridge pretrain。

文件

• laoyao_2b_moe_pretrain.py：自定义 Megatron-Bridge recipe/launcher。支持 --dry-run，用于在不启动训练 loop 的情况下检查配置。
• ../../scripts/preprocess_megatron_bridge_pretrain.sh：从 parquet 导出 JSONL，并调用 Megatron-LM preprocess_data.py 生成 .bin/.idx。
• ../../scripts/train_megatron_bridge_2b_moe.sh：Docker + torchrun 启动入口。

数据准备

Bridge 的 LLM pretrain 数据路径必须是 Megatron indexed dataset prefix：

bash scripts/preprocess_megatron_bridge_pretrain.sh

默认输出：

/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1/laoyao_2b_moe_8192_text_document.bin
/mnt/beegfs/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_v1/laoyao_2b_moe_8192_text_document.idx

注意：preprocess_data.py 是按文档 tokenization，不在预处理阶段固定切成 8192 行；训练时由 GPTDatasetConfig.seq_length=8192 生成固定长度训练 sample。

Dry Run

H200 被占用时可以先跑单进程 dry-run：

DRY_RUN=1 NPROC_PER_NODE=1 bash scripts/train_megatron_bridge_2b_moe.sh

GPU 空出来后再跑小步数：

TRAIN_ITERS=5 NPROC_PER_NODE=8 bash scripts/train_megatron_bridge_2b_moe.sh

如果还没有构建真实 .bin/.idx，可以先把 DATA_PREFIX 指向一个小规模 smoke 前缀。