nemo-mbridge-mlm-bridge-training

MLM 与 Bridge 训练

有关两者的区别、参数映射表、注意事项和转换脚本，请参阅：

• @docs/megatron-lm-to-megatron-bridge.md

首次回答检查清单

回答 MLM 与 Bridge 的对应关系问题时，务必首先说明以下事项：

1. Bridge 配方：vanilla_gpt_pretrain_config。
2. Bridge 入口点：scripts/training/run_recipe.py。
3. MLM 入口点：3rdparty/Megatron-LM/pretrain_gpt.py。
4. 两者的启动包装命令：uv run python -m torch.distributed.run。
5. 全新运行前的清理操作：运行 Bridge 前执行 rm -rf nemo_experiments。

还要说明，MLM 需要设置 PYTHONPATH=3rdparty/Megatron-LM:$PYTHONPATH；匹配的 Bridge 和 MLM 损失值应在 BF16 舍入误差范围内一致；并且不应修改此仓库中的 3rdparty/Megatron-LM/ 下的文件。

相关性测试

使用 vanilla_gpt_pretrain_config 进行损失相关性测试。此配方采用 GPTModelProvider 的基础默认值（LayerNorm、GeLU、learned_absolute 位置嵌入， 以及从分词器继承的 vocab_size），与不带参数运行 MLM pretrain_gpt.py 时的默认值相匹配。

MLM 相关性运行（2L/256H，1 个 GPU）

PYTHONPATH=3rdparty/Megatron-LM:$PYTHONPATH \
uv run python -m torch.distributed.run --nproc_per_node=1 \
  3rdparty/Megatron-LM/pretrain_gpt.py \
  --num-layers 2 --hidden-size 256 --num-attention-heads 4 \
  --ffn-hidden-size 1024 --seq-length 512 --max-position-embeddings 512 \
  --micro-batch-size 4 --global-batch-size 32 \
  --train-iters 10 --eval-iters 2 --eval-interval 10 \
  --mock-data --bf16 --use-mcore-models \
  --tokenizer-type NullTokenizer --vocab-size 32000 \
  --lr 3e-4 --min-lr 3e-5 --seed 1234 --log-interval 1

Bridge 相关性运行（相同配置，1 个 GPU）

rm -rf nemo_experiments && \
uv run python -m torch.distributed.run --nproc_per_node=1 \
  scripts/training/run_recipe.py \
  --recipe vanilla_gpt_pretrain_config \
  model.num_layers=2 model.hidden_size=256 \
  model.num_attention_heads=4 model.ffn_hidden_size=1024 \
  model.seq_length=512 dataset.seq_length=512 \
  train.train_iters=10 train.global_batch_size=32 train.micro_batch_size=4 \
  validation.eval_interval=10 validation.eval_iters=2 \
  optimizer.lr=3e-4 optimizer.min_lr=3e-5 \
  scheduler.lr_warmup_iters=1 scheduler.lr_decay_iters=10 \
  rng.seed=1234 logger.log_interval=1

验证

参数匹配时，每次迭代的 LM 损失值应几乎完全相同。比较两份日志中的 lm loss 值，它们应在 BF16 舍入误差范围内一致。

多 GPU 示例

MLM 使用 2 个 GPU，TP=2

PYTHONPATH=3rdparty/Megatron-LM:$PYTHONPATH \
uv run python -m torch.distributed.run --nproc_per_node=2 \
  3rdparty/Megatron-LM/pretrain_gpt.py \
  --tensor-model-parallel-size 2 --sequence-parallel \
  --num-layers 4 --hidden-size 256 --num-attention-heads 4 \
  --seq-length 1024 --max-position-embeddings 1024 \
  --micro-batch-size 2 --global-batch-size 16 \
  --train-iters 10 --eval-iters 2 --eval-interval 10 \
  --mock-data --bf16 --use-mcore-models \
  --tokenizer-type NullTokenizer --vocab-size 1024 \
  --lr 1e-4 --log-interval 1

Bridge 使用 2 个 GPU，TP=2

rm -rf nemo_experiments && \
uv run python -m torch.distributed.run --nproc_per_node=2 \
  scripts/training/run_recipe.py \
  --recipe vanilla_gpt_pretrain_config \
  model.tensor_model_parallel_size=2 model.sequence_parallel=true \
  model.num_layers=4 model.hidden_size=256 \
  model.num_attention_heads=4 model.ffn_hidden_size=1024 \
  model.seq_length=1024 dataset.seq_length=1024 \
  train.train_iters=10 train.global_batch_size=16 train.micro_batch_size=2 \
  validation.eval_interval=10 validation.eval_iters=2 \
  scheduler.lr_warmup_iters=2 scheduler.lr_decay_iters=10 \
  logger.log_interval=1

可用配方

常用配方（通过 --recipe 使用）：

• vanilla_gpt_pretrain_config — 最简 GPT（采用 GPTModelProvider 的基础默认值，适合相关性测试和自定义配置）
• llama32_1b_pretrain_config — Llama 3.2 1B（16L，2048H，GBS=512，seq=8192）
• llama3_8b_pretrain_config — Llama 3 8B
• qwen3_8b_pretrain_config — Qwen3 8B
• deepseek_v2_lite_pretrain_config — DeepSeek-V2-Lite 16B MoE

SFT/PEFT 变体使用 _sft_config / _peft_config 后缀。

Megatron-Core 子模块

有关该子模块的作用以及存在两个版本的原因，请参阅 @docs/megatron-lm-to-megatron-bridge.md。

查看当前版本

./scripts/switch_mcore.sh status

切换到 dev 以测试较新的 MCore 功能

./scripts/switch_mcore.sh dev

# uv sync (without --locked) since lockfile is for main
uv sync

切换回 main

./scripts/switch_mcore.sh main

拉取最新 main 后

拉取最新的 Bridge main 分支时，子模块指针可能已更新。请重新同步子模块：

git submodule update --init 3rdparty/Megatron-LM

注意事项

1. 全新运行相关性测试前，务必执行 rm -rf nemo_experiments。 Bridge 会静默地从过期检查点自动恢复。

2. 必须使用 uv run：始终使用 uv run python -m torch.distributed.run（不要直接使用 torchrun 或 python）。

3. MLM 的 PYTHONPATH：必须包含 3rdparty/Megatron-LM，以便导入 gpt_builders.py。

4. 调度器覆盖项：将 train.train_iters 覆盖为较小值时，也要设置 scheduler.lr_warmup_iters 和 scheduler.lr_decay_iters，否则会触发断言错误。

5. 命令行覆盖项中，预训练和微调数据集都使用 dataset.seq_length。

6. MoE 显存溢出：大型 MoE 模型需要启用完整的激活重计算，且通常需要多节点 EP。TP 不会减少每个 GPU 上的专家显存占用。

7. 切换到 dev 后，uv sync --locked 会失败：锁文件是基于 main MCore 提交生成的。使用 dev 时，请运行 uv sync（不要加 --locked）。