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.sequence_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.sequence_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. CLI 覆盖中使用 dataset.sequence_length，不要使用 dataset.seq_length。

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

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