distributed-llm-pretraining-torchtitan

TorchTitan - PyTorch ネイティブの分散 LLM 事前学習

クイックスタート

TorchTitan は、組み合わせ可能な 4D 並列化（FSDP2、TP、PP、CP）による大規模 LLM 事前学習のための PyTorch 公式プラットフォームで、H100 GPU 上でベースラインと比べて 65% 以上の高速化を達成しています。

インストール:

# From PyPI (stable)
pip install torchtitan

# From source (latest features, requires PyTorch nightly)
git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt

トークナイザーのダウンロード:

# Get HF token from https://huggingface.co/settings/tokens
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...

8 GPU でのトレーニングを開始:

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh

一般的なワークフロー

ワークフロー 1: 単一ノードでの Llama 3.1 8B の事前学習

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

Single Node Pretraining:
- [ ] Step 1: Download tokenizer
- [ ] Step 2: Configure training
- [ ] Step 3: Launch training
- [ ] Step 4: Monitor and checkpoint

ステップ 1: トークナイザーのダウンロード

python scripts/download_hf_assets.py \
  --repo_id meta-llama/Llama-3.1-8B \
  --assets tokenizer \
  --hf_token=YOUR_HF_TOKEN

ステップ 2: トレーニングの設定

TOML 設定ファイルを編集するか作成します:

# llama3_8b_custom.toml
[job]
dump_folder = "./outputs"
description = "Llama 3.1 8B training"

[model]
name = "llama3"
flavor = "8B"
hf_assets_path = "./assets/hf/Llama-3.1-8B"

[optimizer]
name = "AdamW"
lr = 3e-4

[lr_scheduler]
warmup_steps = 200

[training]
local_batch_size = 2
seq_len = 8192
max_norm = 1.0
steps = 1000
dataset = "c4"

[parallelism]
data_parallel_shard_degree = -1  # Use all GPUs for FSDP

[activation_checkpoint]
mode = "selective"
selective_ac_option = "op"

[checkpoint]
enable = true
folder = "checkpoint"
interval = 500

ステップ 3: トレーニングの起動

# 8 GPUs on single node
CONFIG_FILE="./llama3_8b_custom.toml" ./run_train.sh

# Or explicitly with torchrun
torchrun --nproc_per_node=8 \
  -m torchtitan.train \
  --job.config_file ./llama3_8b_custom.toml

ステップ 4: 監視とチェックポイント

TensorBoard のログは ./outputs/tb/ に保存されます:

tensorboard --logdir ./outputs/tb

ワークフロー 2: SLURM によるマルチノードトレーニング

Multi-Node Training:
- [ ] Step 1: Configure parallelism for scale
- [ ] Step 2: Set up SLURM script
- [ ] Step 3: Submit job
- [ ] Step 4: Resume from checkpoint

ステップ 1: スケールに合わせて並列化を設定

256 GPU（32 ノード）での 70B モデルの場合:

[parallelism]
data_parallel_shard_degree = 32  # FSDP across 32 ranks
tensor_parallel_degree = 8        # TP within node
pipeline_parallel_degree = 1      # No PP for 70B
context_parallel_degree = 1       # Increase for long sequences

ステップ 2: SLURM スクリプトの設定

#!/bin/bash
#SBATCH --job-name=llama70b
#SBATCH --nodes=32
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8

srun torchrun \
  --nnodes=32 \
  --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  -m torchtitan.train \
  --job.config_file ./llama3_70b.toml

ステップ 3: ジョブの送信

sbatch multinode_trainer.slurm

ステップ 4: チェックポイントからの再開

設定されたフォルダーにチェックポイントが存在する場合、トレーニングは自動的に再開されます。

ワークフロー 3: H100 向けの Float8 トレーニングの有効化

Float8 は H100 GPU で 30〜50% の高速化をもたらします。

Float8 Training:
- [ ] Step 1: Install torchao
- [ ] Step 2: Configure Float8
- [ ] Step 3: Launch with compile

ステップ 1: torchao のインストール

USE_CPP=0 pip install git+https://github.com/pytorch/ao.git

ステップ 2: Float8 の設定

TOML 設定に以下を追加します:

[model]
converters = ["quantize.linear.float8"]

[quantize.linear.float8]
enable_fsdp_float8_all_gather = true
precompute_float8_dynamic_scale_for_fsdp = true
filter_fqns = ["output"]  # Exclude output layer

[compile]
enable = true
components = ["model", "loss"]

ステップ 3: compile を有効にして起動

CONFIG_FILE="./llama3_8b.toml" ./run_train.sh \
  --model.converters="quantize.linear.float8" \
  --quantize.linear.float8.enable_fsdp_float8_all_gather \
  --compile.enable

ワークフロー 4: 405B モデル向けの 4D 並列化

4D Parallelism (FSDP + TP + PP + CP):
- [ ] Step 1: Create seed checkpoint
- [ ] Step 2: Configure 4D parallelism
- [ ] Step 3: Launch on 512 GPUs

ステップ 1: シードチェックポイントの作成

PP ステージ間で一貫した初期化を行うために必要です:

NGPU=1 CONFIG_FILE=./llama3_405b.toml ./run_train.sh \
  --checkpoint.enable \
  --checkpoint.create_seed_checkpoint \
  --parallelism.data_parallel_shard_degree 1 \
  --parallelism.tensor_parallel_degree 1 \
  --parallelism.pipeline_parallel_degree 1

ステップ 2: 4D 並列化の設定

[parallelism]
data_parallel_shard_degree = 8   # FSDP
tensor_parallel_degree = 8       # TP within node
pipeline_parallel_degree = 8     # PP across nodes
context_parallel_degree = 1      # CP for long sequences

[training]
local_batch_size = 32
seq_len = 8192

ステップ 3: 512 GPU での起動

# 64 nodes x 8 GPUs = 512 GPUs
srun torchrun --nnodes=64 --nproc_per_node=8 \
  -m torchtitan.train \
  --job.config_file ./llama3_405b.toml

使用すべき場合と代替手段

TorchTitan を使用する場合:

• ゼロから LLM を事前学習する場合（8B から 405B 以上まで）
• サードパーティの依存関係なしで PyTorch ネイティブなソリューションが必要な場合
• 組み合わせ可能な 4D 並列化（FSDP2、TP、PP、CP）が必要な場合
• Float8 サポートを備えた H100 でトレーニングする場合
• torchtune/HuggingFace と相互運用可能なチェックポイントが必要な場合

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

• Megatron-LM: NVIDIA のみの環境で最大のパフォーマンスが必要な場合
• DeepSpeed: より広範な ZeRO 最適化エコシステムや推論サポートが必要な場合
• Axolotl/TRL: 事前学習ではなくファインチューニングを行う場合
• LitGPT: 教育目的や小規模なトレーニングの場合

よくある問題

問題: 大規模モデルでメモリ不足になる

アクティベーションチェックポイントを有効にし、バッチサイズを減らします:

[activation_checkpoint]
mode = "full"  # Instead of "selective"

[training]
local_batch_size = 1

または勾配累積を使用します:

[training]
local_batch_size = 1
global_batch_size = 32  # Accumulates gradients

問題: 非同期コレクティブを使用した TP でメモリ使用量が多い

環境変数を設定します:

export TORCH_NCCL_AVOID_RECORD_STREAMS=1

問題: Float8 トレーニングが速くならない

Float8 は大きな GEMM でのみ効果があります。小さなレイヤーをフィルタリングします:

[quantize.linear.float8]
filter_fqns = ["attention.wk", "attention.wv", "output", "auto_filter_small_kn"]

問題: 並列化の変更後にチェックポイントの読み込みが失敗する

DCP のリシャーディング機能を使用します:

# Convert sharded checkpoint to single file
python -m torch.distributed.checkpoint.format_utils \
  dcp_to_torch checkpoint/step-1000 checkpoint.pt

問題: パイプライン並列化の初期化

先にシードチェックポイントを作成してください（ワークフロー 4、ステップ 1 を参照）。

サポートされているモデル

モデル	サイズ	ステータス
Llama 3.1	8B、70B、405B	Production
Llama 4	各種	Experimental
DeepSeek V3	16B、236B、671B（MoE）	Experimental
GPT-OSS	20B、120B（MoE）	Experimental
Qwen 3	各種	Experimental
Flux	Diffusion	Experimental

パフォーマンスベンチマーク（H100）

モデル	GPU 数	並列化	TPS/GPU	技術
Llama 8B	8	FSDP	5,762	ベースライン
Llama 8B	8	FSDP+compile+FP8	8,532	+48%
Llama 70B	256	FSDP+TP+AsyncTP	876	2D 並列
Llama 405B	512	FSDP+TP+PP	128	3D 並列

高度なトピック

FSDP2 の設定: FSDP2 と FSDP1 の詳細な比較および ZeRO との対応関係については、references/fsdp.md を参照してください。

Float8 トレーニング: テンソルワイズとローワイズのスケーリングレシピについては、references/float8.md を参照してください。

チェックポイント: HuggingFace への変換と非同期チェックポイントについては、references/checkpoint.md を参照してください。

カスタムモデルの追加: TrainSpec プロトコルについては、references/custom-models.md を参照してください。

Resources

• GitHub: https://github.com/pytorch/torchtitan
• Paper: https://arxiv.org/abs/2410.06511
• ICLR 2025: https://iclr.cc/virtual/2025/poster/29620
• PyTorch Forum: https://discuss.pytorch.org/c/distributed/torchtitan/44