MetaのPyTorchチームがTorchTitanの更新を公開し、最新の公開ウェイトモデルと強化学習への対応を追加して、大規模学習の通信・保存の負担を減らした。生成AIを学習させる開発チームは、同じモデル実装で事前学習から強化学習まで進め、複数のGPUで実験規模を拡大できる。
1. モデル別のコード管理の負担を軽減
TorchTitanは、生成AIの事前学習と、事前学習後に目的に合わせて調整する事後学習のためのプラットフォームだ。今回はQwen3.5やDeepSeek-V4などの最新モデルに対応する。モデルの構成要素や並列化方式、ハードウェア別のGPU演算を担うカーネルを変更しても、モデル別コードを個別に維持せずに済む構成の仕組みを提供する。
並列化は、学習作業を複数のGPUに分ける方式だ。開発チームは同じモデルを基に、作業の分割方法や演算部品を置き換えられる。新しいモデルを追加するたびに学習基盤全体を作り直す負担を減らす方向だ。
2. TitanRLで学習と推論のモデル実装を統合
TitanRLは、事前学習・事後学習に使うモデルをvLLM内部の推論にもそのまま使用する。モデル構造の変更時に、学習用と推論用の実装を両方修正する重複作業を減らす。学習用の応答生成は完了を待たずに進める非同期方式で行い、学習と生成をそれぞれ複数の並列化方式で拡張できる。
チームは、入力ごとに一部の計算だけを有効にするエキスパートモジュールに計算を分けるモデルと、全入力に全パラメータを使う密なモデルで試験したと述べた。モデルが一度に参照する入力範囲であるコンテキストやコード実行などを含むタスクで、学習データ処理の反復単位である数百ステップの安定した学習を試した。公開事例の1つは、Qwen3.5-35B-A3BベースモデルをGB300 40台、64Kコンテキストで学習・検証したソフトウェア問題解決実験だ。
チームによると、バッチ不変モード（batch-invariant mode）を使う対応済みのオンポリシー構成では、学習器と生成器の対数確率がビット単位で一致する。これは現在の方策で生成したデータを学習に使う構成で、対数確率は確率の対数を取った値を指す。チームはこれを、数学とターミナルエージェントの強化学習における数値不一致を調べる研究に活用した。
3. 大規模なスループットは実行条件とともに提示
同社発表では、社内GB300クラスターでDeepSeek V3 671Bに次トークン予測を補助するMTP1モジュールを追加した。8ビット浮動小数点形式のMXFP8で数値を低精度形式にする量子化を適用した。機械学習性能の公開評価基準であるMLPerf学習ベンチマークは、GPU 256台構成でGPU当たり毎秒5,324.21トークンを記録した。計算を複数段階に分けるパイプライン段数は2、データを複数GPUに分けて処理するデータ並列度は128、エキスパート計算を分割するエキスパート並列度は64だ。1デバイスで一度に処理するデータのまとまりであるローカルバッチと、バッチを分割した小さなまとまりであるマイクロバッチの数は、それぞれ120だった。
Muse Spark 1.2 OSSの同社発表結果では、96Kコンテキスト・GPU 256台構成のGPU当たりの処理量は、学習済みルーティングで1,500.61、エキスパート間の負荷を強制的に均等化するルーティングで1,839.52トークン/秒だった。192K・512台構成では、それぞれ1,335.86と1,588.22トークン/秒だった。コンテキスト・GPU・並列構成が異なる結果は、同一条件の性能順位として比較できない。
4. 通信削減の鍵はモデルを分割する境界
FlexShardは、モデルの状態をGPU間で分割する境界をより細かく定める。同社発表では、DeepSeek V3をH100 256台で実行する際、複数GPUにまたがって学習状態を更新する分散MuonオプティマイザーにOwned・BlockShard配置を適用した。更新前のAll Gather通信をなくし、従来の完全分割データ並列学習方式であるFSDPのみを使う同等実装に比べ、モデルパラメータを更新するアルゴリズムであるオプティマイザーの区間時間を最大79%削減した。
利点とコストはデータの分割方法に左右される。従来方式では、分割境界と量子化ブロックがずれると、より大きな形式でデータを送るか、圧縮転送に必要な補正値を追加でやり取りする必要があった。新方式は境界を量子化ブロックに合わせることで追加通信をなくし、同社発表ではLlama3 405B・H100 256台で学習ステップ全体の時間を最大20%削減した。
5. チェックポイントの保存待ちも短縮
新しいチェックポイントライブラリは、初回保存前にデータを説明する情報であるメタデータを共有し、その後の保存では複数プロセスが共同でやり取りする集合通信を省く。学習状態を保存するたびに複数プロセスが足並みをそろえる負担を減らす方式だ。同社発表では、H100 512台で以下のように、従来方式より保存待ちと保存・読み込みの時間が短くなった。
表：従来方式と新しいチェックポイント方式の所要時間
条件：H100 GPU 512台、DSv3 671B、チェックポイント間隔50（同社発表による）
単位：トレーナー待機時間はms、保存・読み込み時間はs
| 項目 | DCPv1 平均 | DCPv1 P99 | torch_checkpointing 平均 | torch_checkpointing P99 |
|---|---|---|---|---|
| トレーナー待機時間 | 841.2 | 941.46 | 34.1 | 36.6 |
| チェックポイント保存時間 | 384.4 | 403.1 | 13.1 | 13.2 |
| 再分割なしの読み込み時間 | 194.7 | 198.2 | 79 | 82.3 |
P99は、測定値を小さい順に並べたときの99%地点の値だ。保存による学習の停止時間が減れば、チェックポイントをより頻繁に残す余地が生まれる。導入はTorchTitanリポジトリを複製し、案内ドキュメントに従って始める。
今回の更新の価値は、モデル実装の重複と分散学習の待機時間をともに減らし、新しいモデルの実験を大規模な強化学習につなげる開発経路を整えた点にある。