MLM vs Bridge Training

Pour voir comment ils diffèrent, les tableaux de mappage des arguments, les pièges courants et le script de translation, voir :

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

Correlation Testing

Utilisez vanilla_gpt_pretrain_config pour les tests de corrélation de loss. Cette recipe utilise les defaults bruts de GPTModelProvider (LayerNorm, GeLU, learned_absolute position embeddings, vocab_size hérité du tokenizer) — correspondant aux defaults de MLM pretrain_gpt.py sans arguments.

MLM Correlation Run (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 Correlation Run (même config, 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

Verification

Avec des paramètres appariés, les losses LM doivent être quasi identiques à chaque itération. Comparez les valeurs lm loss des deux logs — elles doivent correspondre à l'arrondi BF16 près.

Multi-GPU Examples

MLM 2-GPU with 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 with 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

Available Recipes

Recipes courantes (utiliser avec --recipe) :

vanilla_gpt_pretrain_config — GPT minimal (defaults bruts de GPTModelProvider, idéal pour les tests de corrélation et les configs personnalisées)
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

Les variantes SFT/PEFT utilisent le suffixe _sft_config / _peft_config.

Megatron-Core Submodule

Pour savoir ce qu'est le submodule et pourquoi deux versions existent, voir @docs/megatron-lm-to-megatron-bridge.md.

Check current version

./scripts/switch_mcore.sh status

Switch to dev for testing newer MCore features

./scripts/switch_mcore.sh dev

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

Switch back to main

./scripts/switch_mcore.sh main

After pulling latest main

Quand vous tirez la dernière branche Bridge main, le pointeur du submodule peut avoir été mis à jour. Re-synchronisez le submodule :

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

Pitfalls

Toujours rm -rf nemo_experiments avant une nouvelle correlation run. Bridge reprend silencieusement depuis des checkpoints obsolètes.
uv run requis : Utilisez toujours uv run python -m torch.distributed.run (pas torchrun ou python seul).
MLM PYTHONPATH : Doit inclure 3rdparty/Megatron-LM pour que gpt_builders.py soit importable.
Scheduler overrides : Quand vous remplacez train.train_iters par une petite valeur, définissez aussi scheduler.lr_warmup_iters et scheduler.lr_decay_iters ou vous obtiendrez une assertion error.
Utilisez dataset.sequence_length dans les overrides CLI, pas dataset.seq_length.
MoE OOM : Les grands modèles MoE nécessitent une recomputation d'activation complète et généralement un EP multi-nœud. TP ne réduit PAS la mémoire expert par GPU.
uv sync --locked échoue après switching to dev : Le lockfile est généré par rapport au commit MCore main. Utilisez uv sync (sans --locked) quand vous êtes sur dev.

mlm-bridge-training