用途:在 Megatron-LM(MLM)与 Megatron Bridge 两个框架中运行训练任务,支持模拟数据(--mock-data)或真实数据。用于对比两者差异、做损失相关性验证,以及多 GPU 训练示例。
主要能力:提供损失相关性测试方法,使用 vanilla_gpt_pretrain_config 配方(bare GPTModelProvider 默认值:LayerNorm、GeLU、学习式绝对位置嵌入、vocab_size 继承 tokenizer),与 MLM pretrain_gpt.py 无参默认配置对齐。可用 --recipe 选择常见配方:vanilla_gpt_pretrain_config(最小 GPT)、llama32_1b_pretrain_config、llama3_8b_pretrain_config、qwen3_8b_pretrain_config、deepseek_v2_lite_pretrain_config;SFT/PEFT 变体使用 _sft_config / _peft_config 后缀。
使用方式:MLM 运行需设置 PYTHONPATH=3rdparty/Megatron-LM:$PYTHONPATH,并用 uv run python -m torch.distributed.run 启动 pretrain_gpt.py。Bridge 运行用 uv run python -m torch.distributed.run 执行 scripts/training/run_recipe.py 并传 --recipe 与 CLI 覆盖参数。多 GPU 示例通过 --nproc_per_node、tensor-model-parallel-size、sequence-parallel 等设置。
验证:参数匹配后,两者每个迭代的 LM loss 应几乎一致,差异应在 BF16 舍入范围内。
限制与注意:相关性运行前必须 rm -rf nemo_experiments,否则 Bridge 会静默从旧检查点自动恢复;必须使用 uv run 而非裸 torchrun;MLM 需包含 3rdparty/Megatron-LM 到 PYTHONPATH;覆盖 train.train_iters 为小值时应同时设置 scheduler.lr_warmup_iters 与 scheduler.lr_decay_iters,否则断言错误;CLI 覆盖用 dataset.sequence_length 而非 dataset.seq_length;大型 MoE 模型需完整激活重计算,通常需多节点 EP,TP 不减少每 GPU 专家内存;切换到 dev 后 uv sync --locked 会失败,应使用 uv sync(不带 --locked)。Megatron-Core 子模块可用 ./scripts/switch_mcore.sh status/dev/main 切换版本。