nemo-mbridge-mlm-bridge-training

MLM vs Bridge Training

For how they differ, the arg mapping tables, gotchas, and translation script, see:

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

First Answer Checklist

For MLM-vs-Bridge correlation questions, always name these items up front:

1. Bridge recipe: vanilla_gpt_pretrain_config.
2. Bridge entry point: scripts/training/run_recipe.py.
3. MLM entry point: 3rdparty/Megatron-LM/pretrain_gpt.py.
4. Launch wrapper for both: uv run python -m torch.distributed.run.
5. Fresh-run cleanup: rm -rf nemo_experiments before the Bridge run.

Also state that MLM needs PYTHONPATH=3rdparty/Megatron-LM:$PYTHONPATH, matched Bridge and MLM losses should agree within BF16 rounding, and files under 3rdparty/Megatron-LM/ should not be modified from this repo.

Correlation Testing

Use vanilla_gpt_pretrain_config for loss-correlation testing. This recipe uses bare GPTModelProvider defaults (LayerNorm, GeLU, learned_absolute position embeddings, vocab_size inherited from tokenizer) — matching MLM pretrain_gpt.py defaults with no args.

MLM Correlation Run (2L/256H, 1 GPU)

PYTHONPATH=3rdparty/Megatron-LM:$PYTHONPATH \
uv run python -m torch.distributed.run --nproc_per_node=1 \

Bridge Correlation Run (same config, 1 GPU)

Verification

With matched parameters the LM losses should be nearly identical at each iteration. Compare lm loss values from both logs — they should agree to within BF16 rounding.

Multi-GPU Examples

MLM 2-GPU with TP=2

Bridge 2-GPU with TP=2

Available Recipes

Common recipes (use with --recipe):

• vanilla_gpt_pretrain_config — Minimal GPT (bare GPTModelProvider defaults, ideal for correlation testing and custom configs)
• 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 variants use _sft_config / _peft_config suffix.

Megatron-Core Submodule

For what the submodule is and why two versions exist, see @docs/megatron-lm-to-megatron-bridge.md.

Check current version

Switch to dev for testing newer MCore features

Switch back to main

After pulling latest main

When you pull the latest Bridge main branch, the submodule pointer may have been updated. Re-sync the submodule:

Pitfalls

1. Always rm -rf nemo_experiments before a fresh correlation run. Bridge auto-resumes from stale checkpoints silently.

2. uv run required: Always use uv run python -m torch.distributed.run (not bare torchrun or python).

3. MLM PYTHONPATH: Must include 3rdparty/Megatron-LM so gpt_builders.py is importable.

4. Scheduler overrides: When overriding train.train_iters to a small value, also set scheduler.lr_warmup_iters and scheduler.lr_decay_iters or you get an assertion error.

5. Use dataset.seq_length in CLI overrides for both pretraining and fine-tuning datasets.

6. MoE OOM: Large MoE models require full activation recomputation and typically multi-node EP. TP does NOT reduce per-GPU expert memory.

7. uv sync --locked fails after switching to dev: The lockfile is generated against the main MCore commit. Use uv sync (without --locked) when on dev.