pytorch/torchtitan

torchtitan – 大規模生成AIモデルの学習を目的としたPyTorchネイティブプラットフォーム

何であるか – torchtitan はPyTorchチームによって開発されたオープンソースライブラリであり、Llama 3.1などの大規模言語モデル（LLM）をネイティブPyTorchプリミティブのみを使用してスケーラブルに学習できるようにします。FSDP2、テンソル並列、パイプライン並列、コンテキスト並列など、分散学習技術のコレクションを、チェックポイント、量子化、ログ記録のユーティリティと統合し、洗練された拡張可能なPython APIを通じて提供します。

なぜ重要か – 現代のLLM（8 B～405 Bパラメータ）を学習するには、カスタムCUDAカーネル、サードパーティライブラリ、重いオーケストレーションツールの複雑なスタックが必要です。torchtitanは、最新のPyTorch nightlyリリースと直接連携する最小限のコードベースを提供することで、この複雑さを軽減することを目指しています。研究者はモデルコードを再実装せずに新しい並列化戦略を実験でき、プロダクションチームは同じライブラリでプロトタイプからマルチノード学習へ移行できます。

コア機能（READMEに記載）

カテゴリ	機能
並列処理	• 多次元で組み合わせ可能な並列処理：FSDP2（パラメータごとのシャーディング）、テンソル並列（非同期TPを含む）、パイプライン並列（ゼロバブル）、超長シーケンス向けのコンテキスト並列。
メモリと速度の最適化	• メタデバイスモデル初期化、選択的／完全なアクティベーションチェックポイント、分散非同期チェックポイント、BF16オプティマイザステート、Float8およびMXFP8量子化（Blackwell GPU対応）。
学習ワークフロー	• `torch.compile`対応、トークン数フラグによる勾配蓄積、柔軟なウォームアップ・安定・減衰型LRスケジューラ、C4を事前構成済みの組み込みデータローダ、カスタムデータセットプラグイン。
可観測性	• TensorBoard / Weights & Biasesメトリクス（損失、メモリ、スループット、MFU）、ランクごとの構造化ログ、プロファイリングツール（CPU/GPU、フライトレコーダ）。
相互運用性	• 分散チェックポイント形式（DCP）は`torchtune`で直接読み込み可能で、ファインチューニングに使用可能。TorchFT統合、トークナイザダウンロードおよびチェックポイント変換用のヘルパースクリプト。
拡張性	• よくドキュメント化された拡張ポイント、迅速なプロトタイピング用の`experiments`フォルダ（例：TitanRL RLスタック）、新しいモデルを追加するための明確なガイドライン。
対応ハードウェア	• NVIDIA GPU（H100までテスト済み、512GPU実行可能）、AMD ROCmフォーク、任意のCUDA/ROCmバージョンに対応するnightlyビルド。

一般的なワークフロー

インストール – pip install torchtitan（またはnightly PyTorch + nightly torchtitanのウェHEELを使用）。単一ノードまたはSlurm/ParallelClusterクラスタ上で即座に動作します。
データ準備 – 組み込みのC4ローダを使用するか、カスタムデータセットを設定ファイルで指定。
モデル選択 – Llama 3.1（8 B、70 B、405 B）が同梱。新しいモデルを追加するにはtorchtitan/models/README.mdの短いガイドに従います。
設定 – Python設定ファイルとCLIフラグ（--module, --config）ですべてのオプションを制御可能。並列処理、チェックポイント、量子化の有効／無効を自由に切り替えられます。
起動 – 1〜Nノードでtorchrun（または提供されるrun_train.sh）を実行。torchtitanは要求された並列処理を自動的に組み合わせ、学習を開始します。
監視 – メトリクスはTensorBoard/W&Bに表示。ログにはデバッグ用のランクごとのトレースが含まれます。
チェックポイントとファインチューニング – 保存されたDCPチェックポイントはtorchtuneで直接読み込み、下流のSFTやRLHFに使用可能。

クイックスタート例

# 1️⃣ nightly PyTorch + torchtitanのウェHEELをインストール
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130
pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130

# 2️⃣ Llama‑3.1トークナイザをダウンロード
python scripts/download_hf_assets.py \
    --repo_id meta-llama/Llama-3.1-8B \
    --assets tokenizer \
    --hf_token <your‑HF‑token>

# 3️⃣ 8GPU実行（単一ノード）
MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

CONFIG=llama3_8bファイルはtorchtitan/config/にあり、既にFSDP2 + テンソル並列 + アクティベーションチェックポイントなどを有効化しています。

どの人にとって適しているか？

研究者：新しいスケーリング法則、並列化の組み合わせ、量子化技術を探索する人。クリーンルーム実装により、内部コードの読解と修正が容易です。
エンジニア：プロダクション品質のLLM事前学習パイプラインを構築したい人。外部のプロプライエタリスタックを避け、PyTorchエコシステム内に留まりたい人。
学生：分散学習の概念を学びたい人。experimentsフォルダには実行可能なデモ（例：TitanRL）が含まれており、同じモデルコードがRL学習や推論に再利用できる様子が示されています。

コミュニティとサポート

フォーラム – 専用のPyTorchフォーラムカテゴリ（pytorch‑forum → distributed → torchtitan）。
論文 – ICLR 2025論文 「TorchTitan: プロダクション対応LLM事前学習のためのワンストップPyTorchネイティブソリューション」（arXiv 2410.06511）。
貢献 – コア変更や実験的アイデアの追加に関するガイドラインはCONTRIBUTING.mdおよびtorchtitan/experiments/README.mdに記載。
サードパーティフォーク – AMD最適化版フォークが存在（AMD‑AGI/torchtitan‑amd）。

ライセンス

torchtitanはBSD‑3‑Clauseライセンスの下でリリースされています。コード自体は許容的なライセンスですが、ダウンロードする外部データやモデル重みのライセンスに従う必要があります（例：Meta‑Llamaトークナイザおよびチェックポイント）。

結論 – torchtitanは、生産環境対応のPyTorchのみのスタックを提供し、大規模LLMの学習を可能にします。読みやすさ、拡張性、最新の分散学習研究に重点を置いています。重い外部エコシステムを導入せずに大規模生成モデルの学習や実験が必要な場合、このライブラリから始めるべきです。