laoyao_2b_moe/training/megatron_bridge

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。
• 训练数据 blend 优先读取 prefix_category_stats.json，按目标 category mix 计算 shard 权重；如果没有 category stats，则退回按 manifest token 数加权。
• validation 使用单独 heldout 2.8k Megatron indexed prefix，不再从训练 split 里切 valid。

文件

• laoyao_2b_moe_pretrain.py：自定义 Megatron-Bridge recipe/launcher。包含 GLM-5.2 vocab、MoE grouped GEMM、heldout validation、checkpoint 保留、recompute、distributed optimizer/overlap，以及信号触发的临时 checkpoint。
• ../../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。

当前 g0050 主训练使用直接 parquet -> Megatron indexed dataset 的路径：

SOURCE_DATA_DIRS=/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_200b_v1_20260701:/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/train/logic_topup_proof_pile_17b_v1_20260701 \
WORK_DIR=/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_glm52_direct_v1 \
MAX_SEQ_LEN=8192 \
bash scripts/preprocess_megatron_bridge_pretrain_direct.sh

该转换脚本会在 manifest.json 中写入可直接传给 Megatron 的 _text_document prefix；不要手动拼接错误的 prefix。

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 前缀。

Graceful Exit

训练入口默认设置：

cfg.train.exit_signal_handler = True
cfg.train.exit_signal = signal.SIGINT

这使用 Megatron-Bridge 自己的 checkpoint_and_decide_exit 路径。signal handler 只记录退出请求；标准训练循环会在 iteration 边界执行 checkpoint，因此不会从 Python signal callback 中直接运行 CUDA、NCCL 或文件写入。

训练 launcher 会将活动训练信息写入 .runtime/active_training.txt。操作时直接运行 scripts/graceful_stop_megatron_training.sh，无需查询 container name；不要把 SIGINT 直接发给 torchrun/container PID 1。