training-llms-megatron

Megatron-Core - 大規模 LLM トレーニング

クイックスタート

Megatron-Core は、高度な並列化戦略により、2B から 462B パラメータの LLM を H100 GPU 上で最大 47% のモデル FLOP 利用率（MFU）でトレーニングします。

インストール:

# Docker (recommended)
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.04-py3

# Or pip
pip install megatron-core

シンプルな分散トレーニング:

# Train with 2 GPUs using data parallelism
torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py

# Or LLaMA-3 8B training
./examples/llama/train_llama3_8b_fp8.sh

一般的なワークフロー

ワークフロー 1: 3D 並列化による LLaMA 形式モデルのトレーニング

このチェックリストをコピーしてください:

LLaMA Training Setup:
- [ ] Step 1: Choose parallelism configuration
- [ ] Step 2: Configure training hyperparameters
- [ ] Step 3: Launch distributed training
- [ ] Step 4: Monitor performance metrics

ステップ 1: 並列化構成を選択する

モデルサイズに応じて並列化戦略を決めます:

モデルサイズ	GPU 数	テンソル並列	パイプライン並列	データ並列	コンテキスト並列
7B	8	1	1	8	1
13B	8	2	1	4	1
70B	64	4	4	4	1
405B	128	8	8	2	2

ステップ 2: トレーニングのハイパーパラメーターを設定する

#!/bin/bash
# train_llama_70b.sh

GPUS_PER_NODE=8
NNODES=8  # 64 GPUs total
TP=4      # Tensor parallel
PP=4      # Pipeline parallel
CP=1      # Context parallel

# LLaMA 70B configuration
MODEL_SIZE=70  # Billion parameters
HIDDEN_SIZE=8192
NUM_LAYERS=80
NUM_HEADS=64
SEQ_LENGTH=4096

# Training hyperparameters
MICRO_BATCH=1
GLOBAL_BATCH=1024
LR=3e-4

torchrun \
  --nproc_per_node=$GPUS_PER_NODE \
  --nnodes=$NNODES \
  pretrain_gpt.py \
  --tensor-model-parallel-size $TP \
  --pipeline-model-parallel-size $PP \
  --context-parallel-size $CP \
  --sequence-parallel \
  --num-layers $NUM_LAYERS \
  --hidden-size $HIDDEN_SIZE \
  --num-attention-heads $NUM_HEADS \
  --seq-length $SEQ_LENGTH \
  --max-position-embeddings $SEQ_LENGTH \
  --micro-batch-size $MICRO_BATCH \
  --global-batch-size $GLOBAL_BATCH \
  --lr $LR \
  --train-iters 100000 \
  --lr-decay-style cosine \
  --lr-warmup-iters 2000 \
  --weight-decay 0.1 \
  --clip-grad 1.0 \
  --bf16 \
  --use-mcore-models \
  --transformer-impl transformer_engine \
  --data-path /path/to/data \
  --vocab-file /path/to/vocab.json \
  --merge-file /path/to/merges.txt

ステップ 3: 分散トレーニングを起動する

# Single node (8 GPUs)
bash train_llama_70b.sh

# Multi-node with SLURM
sbatch --nodes=8 --gpus-per-node=8 train_llama_70b.sh

ステップ 4: パフォーマンスメトリクスを監視する

追跡すべき主要なメトリクス:

Model FLOP Utilization (MFU): Target >40% on H100
Throughput: Tokens/sec/GPU
Memory usage: <80GB per GPU for 70B model
Loss: Should decrease steadily

ワークフロー 2: Mixture of Experts（MoE）トレーニングの構成

Mixtral のようなスパース MoE モデル向けです。

MoE Training:
- [ ] Step 1: Configure expert parallelism
- [ ] Step 2: Set MoE hyperparameters
- [ ] Step 3: Launch training with EP

ステップ 1: エキスパート並列を構成する

# Mixtral 8x7B example
TENSOR_PARALLEL=2
PIPELINE_PARALLEL=1
EXPERT_PARALLEL=4  # Split 8 experts across 4 GPUs
DATA_PARALLEL=4

TOTAL_GPUS=$((TENSOR_PARALLEL * PIPELINE_PARALLEL * EXPERT_PARALLEL * DATA_PARALLEL))
# = 2 * 1 * 4 * 4 = 32 GPUs

ステップ 2: MoE のハイパーパラメーターを設定する

torchrun \
  --nproc_per_node=8 \
  pretrain_gpt.py \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 1 \
  --expert-model-parallel-size 4 \
  --num-experts 8 \
  --moe-router-topk 2 \
  --moe-router-load-balancing-type aux_loss \
  --moe-aux-loss-coeff 0.01 \
  --hidden-size 4096 \
  --num-layers 32 \
  --num-attention-heads 32 \
  --seq-length 4096 \
  --max-position-embeddings 4096 \
  --bf16 \
  --use-mcore-models \
  --transformer-impl transformer_engine \
  --data-path /path/to/data \
  --vocab-file /path/to/vocab.json \
  --merge-file /path/to/merges.txt

ステップ 3: EP を使ってトレーニングを起動する

エキスパート並列は、異なるエキスパートを GPU 間に分散させるため、キャパシティを維持しつつメモリ使用量を削減できます。

Memory without EP: 8 experts × 7B = 56GB per GPU
Memory with EP=4: 2 experts × 7B = 14GB per GPU
Savings: 75% memory reduction

ワークフロー 3: スループットの最大化

H100 で 47% の MFU を達成します。

Performance Optimization:
- [ ] Step 1: Enable Flash Attention
- [ ] Step 2: Use FP8 precision (H100)
- [ ] Step 3: Optimize micro-batch size
- [ ] Step 4: Tune parallelism degrees

ステップ 1: 最適化を有効にする

--use-mcore-models  # Use Megatron Core models
--transformer-impl transformer_engine  # Use Transformer Engine
--sequence-parallel  # Reduce activation memory (use with TP)

ステップ 2: FP8 精度を使用する（H100 のみ）

--fp8-hybrid  # FP8 mixed precision training
# Transformer Engine handles FP8 automatically

結果: BF16 と比べて、H100 で 1.5〜2 倍の高速化が得られます。

ステップ 3: マイクロバッチサイズを最適化する

メモリに収まる最大のマイクロバッチを見つけます:

# Start with 1, increase until OOM
for MBS in 1 2 4 8; do
  echo "Testing micro-batch-size=$MBS"
  torchrun ... --micro-batch-size $MBS
done

代表的な値:

• 7B モデル: 4〜8
• 70B モデル: 1〜2
• 405B モデル: 1

ステップ 4: 並列度を調整する

経験則:

Tensor Parallel: Use ≤8 (limited by NVLink within node)
Pipeline Parallel: Use for >70B models
Context Parallel: Use for sequences >8K tokens
Data Parallel: Fill remaining GPUs

128 台の H100 で 405B モデルを実行する例:

TP=8 (1 node)
PP=8 (across nodes)
CP=2 (long sequences)
DP=1
Total = 8 × 8 × 2 × 1 = 128 GPUs

使用場面と代替手段

Megatron-Core を使用する場合:

• 10B を超えるパラメーターのモデルをトレーニングするとき
• 最大限の効率が必要なとき（目標 40% を超える MFU）
• NVIDIA GPU（A100、H100）を使用しているとき
• 大規模な本番トレーニングを行うとき
• きめ細かな並列化制御が必要なとき

代わりに以下を使用してください:

• PyTorch FSDP: 70B 未満のモデル、よりシンプルな API、PyTorch ネイティブ
• DeepSpeed: セットアップが容易、100B 未満のモデルに適している
• HuggingFace Accelerate: プロトタイピング、よりシンプルなワークフロー
• LitGPT: 教育用途、単一ファイルの実装

よくある問題

問題: GPU 使用率が低い（MFU が 30% 未満）

原因:

1. マイクロバッチが小さすぎる
2. 並列化のオーバーヘッドが大きすぎる
3. Flash Attention を使用していない

対処法:

# Increase micro-batch
--micro-batch-size 4  # Was 1

# Enable optimizations
--use-flash-attn
--sequence-parallel

# Reduce TP if >8
--tensor-model-parallel-size 4  # Was 16

問題: メモリ不足（Out of memory）

次の方法でメモリを削減します:

--tensor-model-parallel-size 2  # Split model across GPUs
--recompute-granularity full  # Gradient checkpointing
--recompute-method block  # Checkpoint transformer blocks
--recompute-num-layers 1  # Checkpoint every layer

または CPU/NVMe オフロードを使用します:

--cpu-optimizer  # Offload optimizer to CPU
--cpu-optimizer-type ADAM  # CPU Adam variant

問題: トレーニングが想定より遅い

確認事項:

1. ネットワークのボトルネック: InfiniBand/NVLink が有効になっていることを確認する
2. パイプラインバブル: インターリーブ型パイプラインスケジュールを使用する

   --num-layers-per-virtual-pipeline-stage 2
   

3. データ読み込み: 高速なデータローダーを使用する

   --dataloader-type cyclic
   

問題: 損失が発散する

トレーニングを安定させるには:

--lr-warmup-iters 2000  # Longer warmup
--clip-grad 1.0  # Gradient clipping
--init-method-std 0.006  # Smaller init
--attention-dropout 0.0  # No dropout in attention
--hidden-dropout 0.0  # No dropout in FFN

応用トピック

並列化戦略: TP/PP/DP/CP/EP の詳細な比較とパフォーマンス分析、および各手法の使いどころについては、references/parallelism-guide.md を参照してください。

パフォーマンスベンチマーク: さまざまなモデルサイズと GPU 構成における MFU の数値については、references/benchmarks.md を参照してください。

本番環境の構成: LLaMA 3 405B、Nemotron-4 340B、DeepSeek-V3 671B の実際の構成例については、references/production-examples.md を参照してください。

トレーニングレシピ: GPT/LLaMA/Mixtral アーキテクチャ向けの完全なハイパーパラメーター構成については、references/training-recipes.md を参照してください。

ハードウェア要件

• GPU: NVIDIA Ampere 以降（A100、H100、B200）
  • Turing でも動作するが、速度は遅い
  • FP8 には Hopper/Ada/Blackwell が必要
• ネットワーク: マルチノードには InfiniBand または 400Gb 以上の Ethernet
• GPU あたりのメモリ:
  • 7B モデル: 40GB 以上
  • 70B モデル: 80GB（TP=4 の場合）
  • 405B モデル: 80GB（TP=8、PP=8 の場合）
• ストレージ: チェックポイント用の高速 NVMe（70B 以上のモデルでは 1TB 以上）

リソース

• ドキュメント：https://docs.nvidia.com/megatron-core/
• GitHub：https://github.com/NVIDIA/Megatron-LM
• 論文：
  • 「Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism」（2019年）
  • 「Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM」（2021年）
• NeMo Framework：https://docs.nvidia.com/nemo-framework/（Megatron-Core を基盤としています）

エージェントの作業手順

1. 環境を確認します。 GPU の種類、メモリ、ドライバー／CUDA のバージョン（nvidia-smi）、フレームワークのバージョン、データセットの場所とサイズを確認します。
2. 入力を確定します。 データまたはユーザーに確認して、形式、識別子、パラメーターを確定します。結果に影響する値は推測せず、質問してください。
3. まず小規模に実行します。 スモークテストとして、小さなモデルまたはサブセットを少ないステップ数で実行し、損失が減少することとチェックポイントが保存されることを確認します。
4. 上記の手順と参考情報を使用して、タスク全体を実行します。
5. 結果を検証します。 保留データのメトリクスを追跡し、ベースラインと比較して、シード、設定、ハードウェアを記録します。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、未確定の点を明記します。

発生した問題	対応方法
CUDA のメモリ不足	バッチサイズを小さくする、勾配蓄積／チェックポイントまたは混合精度を有効にする、あるいはモデルをシャーディングします。
損失が NaN になる、または発散する	学習率を下げ、データに無効な値がないか確認し、勾配クリッピングを有効にします。
これらの手順にある関数、フラグ、エンドポイントがインストール済みバージョンに存在しない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認して適宜変更し、ユーザーに伝えます。API を捏造してはいけません。
必須の入力、識別子、パラメーターが曖昧	ユーザーに質問するか、実行前に仮定を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計情報を捏造してはいけません。実行または検証できない場合は、その旨を明確に伝えます。
• ログなしに学習結果を主張してはいけません。大規模なジョブを開始する前に計算コストを見積もります。
• ここに記載されたバージョン固有の詳細は古くなっている可能性があるため、インストール済みバージョンの公式ドキュメントで確認します。
• 費用が発生する、共有 GPU やクラウドクォータを消費する、個人データまたは患者データに触れる、あるいは元に戻せない操作を行う前には、必ず確認を取ります。

関連スキル

• tensorrt-llm：NVIDIA TensorRT で LLM 推論を最適化し、スループットを最大化してレイテンシを最小化します。
• moe-training：DeepSpeed または HuggingFace を使用して、Mixture of Experts（MoE）モデルを学習します。
• axolotl：Axolotl による大規模言語モデルのファインチューニングについて、YAML 学習設定、LoRA、QLoRA、選好学習などを網羅したガイダンスを提供します…