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 の例

TP=2 の MLM、2 GPU

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

TP=2 の Bridge、2 GPU

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 サブモジュール

サブモジュールの概要と、2 つのバージョンが存在する理由については、 @docs/megatron-lm-to-megatron-bridge.md を参照してください。

現在のバージョンを確認する

./scripts/switch_mcore.sh status

新しい MCore 機能をテストするために dev に切り替える

./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：gpt_builders.py をインポートできるように、3rdparty/Megatron-LM を含める必要があります。

4. スケジューラーのオーバーライド：train.train_iters を小さい値にオーバーライドする場合は、scheduler.lr_warmup_iters と scheduler.lr_decay_iters も設定してください。設定しないとアサーションエラーが発生します。

5. CLI のオーバーライドでは dataset.seq_length ではなく、dataset.sequence_length を使用してください。

6. MoE の OOM：大規模な MoE モデルでは、完全なアクティベーション再計算と、通常はマルチノード EP が必要です。TP では GPU あたりのエキスパートメモリーは削減されません。

7. dev に切り替えた後に uv sync --locked が失敗する場合：ロックファイルは main の MCore コミットを基に生成されています。dev 使用時は uv sync（--locked なし）を使用してください。