torchtitan – PyTorchネイティブの大規模生成AIトレーニングプラットフォーム

概要

torchtitan は、PyTorchチームによるオープンソースライブラリで、大規模言語モデル（LLM）およびマルチモーダルモデルを完全にPyTorchでトレーニング、ファインチューニング、強化学習（RL）を行うことができます。コードベースは意図的に小さく読みやすく設計されており、トレーニングパイプライン全体をPythonの「レシピ」、つまり設定オブジェクトを返す関数として表現します。レシピを変更することで、モデルサイズ、データ、並列化戦略を変更したり、モデルを書き直すことなく事前学習からRLに切り替えることもできます。

コアコンセプトとその重要性

アイデア	機能	有用性
構成可能な並列化	モデルはシャーディング方法を宣言し、ライブラリは`DeviceMesh`と`DTensor`を介してFully-Sharded Data Parallel（FSDP2）、Tensor Parallel（TP）、Pipeline Parallel（PP）、Context Parallel（CP）、Expert Parallel（EP）をスタックします。	同じコード上で任意の並列化の組み合わせを試すことができ、カスタムカーネルなしで単一GPUから数百GPUまでスケールします。
事前学習からRLまでのワンスタック	同じモデル定義とカーネルが、古典的な言語モデルの事前学習、教師ありファインチューニング（SFT）、およびvLLM生成を強化するTitanRL強化学習ループにも使用されます。	トレーナーと推論の間でのビット単位の再現性を保証し、RLベースのアライメントやポリシー勾配のデバッグを簡素化します。
最先端モデルのサポート	Llama 3、Qwen 3、DeepSeek V3、FLUX、Kimi、Muse Glimmerなどの既製レシピがあり、低精度フォーマット（MXFP8、NVFP4）とMoE（Dist-MoE）も含まれます。	最先端のモデルでアーキテクチャや量子化のトリックを再実装することなく実験を開始できます。
Pythonファーストの設定	レシピは設定オブジェクトを返すPython関数にすぎず、長いCLIフラグのリストを渡す代わりに関数を編集します。	実験を再現可能でバージョン管理しやすくし、レシピファイルを同僚と共有すれば、完全に同じ実行結果を得られます。
高度な可観測性	組み込みのメトリクス（損失、TFLOPs、MFU）、TensorBoard/W&Bロギング、プロファイリングツール、決定論モード、損失比較スクリプト、SDCリプレイ。	数百GPU規模で重要なパフォーマンスのボトルネックや数値バグを早期に発見できます。

現在できること

新しいLLMをトレーニング – llama3_8bなどのレシピを選択し、8GPUで./run_train.shを実行します。ライブラリはシャーディング、チェックポイント、データロードを自動的に処理します。
LoRAまたはSFTでのファインチューニング – 同じレシピシステムでLoRAアダプターやチャットデータセットのファインチューニング手順を追加できます。
RLベースのアライメント – TitanRLスタックを使用してPPOスタイルやより新しいGRPO/DAPO損失を実行し、トレーニングされた正確なモデル重みを再利用します。
低精度トレーニングの実験 – 設定フラグでMXFP8またはNVFP4を有効にします。ライブラリは必要なtorchaoカーネルを同梱しています。
数百GPUへのスケール – ドキュメント化されたベンチマークでは、H100上で最大512GPUの実行と、自動分散チェックポイントと非同期保存が示されています。

誰が使うべきか

研究ラボ – 新しい並列化やRLのアイデアを探求するためのクリーンでハッキング可能なコードベースが必要な場合。
エンジニア – PyTorchエコシステム内に留まりたい（別のC++/CUDAフレームワークは不要）本番グレードのLLMトレーニングパイプラインを構築している場合。
学生 – 大規模トレーニングを学ぶための、FSDP、TP、MoEの組み合わせ方を示す読みやすいレシピを提供。

はじめに（クイックサマリー）

インストール – リポジトリをクローンし、PyTorchのナイトリーとtorchtitanパッケージをインストールします（またはビルド済みのpip/condaホイールを使用）。
デバッグジョブを実行 – NGPU=1 ./run_train.shで小さなLlama-3デバッグモデルを10ステップトレーニングします。
実際のモデルをトレーニング – Hugging Faceからトークナイザーをダウンロードし、MODULE=torchtitan_recipes.models.llama3 CONFIG=llama3_8bを設定し、8GPUで起動します。
カスタマイズ – 並列化度、データミックス、LoRA追加などを調整する独自のレシピ関数を作成します。

ドキュメントとコミュニティ

ドキュメントは、構成、モデル、分散、データ、チェックポイント、RL、デバッグ、ベンチマークに分かれています。各フォルダには使用例を含むREADME.mdがあります。
CIバッジは、CPU、NVIDIA（H100/B200を含む）、AMD ROCm、およびRLワークロードでのアクティブなテストを示しています。
貢献 – ガイドラインは、experiments/フォルダでの新しい実験と、プルリクエストによるコア改善を推奨しています。
論文と引用 – ICLR-2025の論文で設計選択の詳細が説明されています。torchtitanを基にした研究成果を発表する際は引用してください。

上記の情報はすべてリポジトリのREADMEから直接取得したものです。外部の前提は追加されていません。

pytorch/torchtitan