nemo-mbridge-mlm-bridge-training
Run Megatron-LM (MLM) and Megatron Bridge training with mock or real data. Covers correlation testing, available recipes, and multi-GPU examples.
By nvidia · 1,802 installs
npx skills add nvidia/skills --skill nemo-mbridge-mlm-bridge-training
Source repository · Upstream listing
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)
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.