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.