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 を参照してください。

リソース

• GitHub: https://github.com/pytorch/torchtitan
• 論文: 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

エージェントの運用手順

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

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

インテグリティに関するルール

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

関連スキル

• torchforge-rl-training：インフラストラクチャとアルゴリズムを分離する Meta のライブラリ torchforge を使った、PyTorch ネイティブなエージェント型 RL に関するガイダンスを提供します。
• pytorch-fsdp2：初期化、シャーディング、混合精度／オフロード設定、分散チェックポイン…
• dask：メモリに収まらない規模の pandas／NumPy ワークフロー向けの分散コンピューティング。