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.seq_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.seq_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 を使用してください。

6. MoE の OOM：大規模な MoE モデルでは、完全な活性化再計算と、通常はマルチノードの EP が必要です。TP を使用しても、GPU あたりのエキスパートメモリーは削減されません。

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