TorchTitan - PyTorch ネイティブ LLM プリトレーニング チートシート
TorchTitan はオープンソースの PyTorch ネイティブ大規模生成モデル（LLM）プリトレーニング向けプラットフォーム。PyTorch 最新分散トレーニング技術のクリーンルーム実装で、4D 並列化 — データ・テンソル・パイプライン・コンテキスト並列 — をモジュール方式で構成。FSDP2、torch.compile、分散チェックポイント、エラスティックスケーリングの上に実装。PyTorch ネイティブスケーリングを使いたいチーム向けのリファレンス実装。
分散プリトレーニングはリソース集約的。大規模マルチノード実行前に小さいモデル/スケールで設定を検証。
インストール
| 方法 | コマンド |
|---|
| ソースから | git clone https://github.com/pytorch/torchtitan && cd torchtitan && pip install -e . |
| 要件 | 最新 PyTorch（nightly 推奨）、CUDA GPU |
| Datasets | HF datasets / tokenizers を使用 |
| 確認 | 1 GPU でデバッグ設定を実行 |
並列化の次元（4D）
| 次元 | 分割対象 |
|---|
| Data Parallel（FSDP2） | バッチ、パラメータ/オプティマイザー/勾配を分割 |
| Tensor Parallel | 個別レイヤーを GPU 間に分割 |
| Pipeline Parallel | モデルを GPU 間のステージに分割 |
| Context Parallel | 長いシーケンスを GPU 間に分割 |
これらは構成：例えば dp × tp × pp × cp でモデルを大型 GPU メッシュにマップ。
設定
TorchTitan は TOML 設定ファイル + コマンドラインオーバーライドで駆動。
# 設定でトレーニング起動（torchrun 経由マルチ GPU）
CONFIG_FILE=./train_configs/llama3_8b.toml ./run_train.sh
| 設定 section | 制御内容 |
|---|
[model] | アーキテクチャ、サイズ、フレーバー |
[training] | バッチサイズ、ステップ、シーケンス長、LR |
[parallelism] | data_parallel_shard_degree 、 tensor_parallel_degree 、 pipeline_parallel_degree 、 context_parallel_degree |
[optimizer] | オプティマイザー + ハイパーパラメーター |
[checkpoint] | 分散チェックポイント頻度/パス |
[activation_checkpoint] | 再計算ポリシー |
キー フィーチャー
| フィーチャー | 利点 |
|---|
| FSDP2 | モダン fully-sharded data parallel |
torch.compile | グラフコンパイル高速化 |
| Async 分散チェックポイント | 巨大実行のセーブ/リジューム効率化 |
| Float8 / 低精度 | サポートハードウェアでの高速トレーニング |
| Elastic scaling | 変動 GPU 可用性に適応 |
| Meta-device init | OOM なしで巨大モデル構築 |
マルチノード起動
# ノード間での torchrun（概念） |
torchrun --nnodes=4 --nproc_per_node=8 \
--rdzv_backend=c10d --rdzv_endpoint=$MASTER:29500 \
-m torchtitan.train --job.config_file llama3_70b.toml
| オーバーライド | 例 |
|---|
| 並列度 | --parallelism.tensor_parallel_degree 8 |
| ステップ | --training.steps 10000 |
| Seq length | --training.seq_len 8192 |
| Compile | --training.compile |
監視・チェックポイント
| 項目 | 実行方法 |
|---|
| Logging | TensorBoard / W&B 統合 |
| メトリクス | Loss、tokens/sec、MFU（model FLOPs utilization） |
| Checkpoints | 分散チェックポイント（DCP）セーブ/ロード |
| Resume | チェックポイントステップに設定をポイント |
TorchTitan vs その他トレーニングフレームワーク
| 項目 | TorchTitan | Megatron-LM | Nanotron |
|---|
| Origin | PyTorch team | NVIDIA | Hugging Face |
| Style | PyTorch ネイティブ、構成可能 | 機能豊富、成熟 | ミニマリズム 3D |
| 並列化 | 4D（dp/tp/pp/cp） | 3D+（広範） | 3D |
| 最適用途 | PyTorch ネイティブリファレンススケーリング | max スケール MoE/dense | 小型、ハッキング可能 3D |
リソース
