torchtune: PyTorch native post-training library

本論文は、大規模言語モデルの事後学習ライフサイクルを効率化するために、モジュール性、透明性、拡張性を優先して設計されたPyTorchネイティブのライブラリであるtorchtuneを紹介し、競争力のある性能とメモリ効率を維持しつつ、効率的なファインチューニングと迅速な研究イテレーションを可能にする。

要約

巨大で非常に賢いロボット（大規模言語モデル）が、膨大な量の書籍から読み書きをすでに学習したと想像してください。ここで、詩を書くことや医療質問に答えることなど、特定の新しいスキルを教えたくなります。このプロセスは「ポストトレーニング」または「ファインチューニング」と呼ばれます。

本論文では、この教育プロセスをより速く、安価に、かつ理解しやすくするための新しいツールキット「torchtune」を紹介しています。その仕組みを、簡単なアナロジーを用いて以下に説明します。

1. 課題：「ブラックボックス」対「レゴセット」

torchtune 以前、これらのロボットを教育するためのツールのほとんどは、組み立て済みの家具のようでした。テーブル（トレーニングレシピ）を購入すれば、それは非常に良く機能しましたが、脚や仕上げを変更したい場合、それをハンマーで叩き壊す必要がありました。これらのツールはしばしば、他の巨大で複雑なシステムの上に構築されていたため、修正や微調整が困難でした。何か問題が起きた場合、指示が他のソフトウェアの層の奥に隠れているため、なぜそうなったのかを把握できませんでした。

torchtune は異なります。それはレゴセットのようです。

• モジュール性: 巨大なブロック一つではなく、個々のレンガ（モデルビルダー、データローダー、オプティマイザー）を提供します。全体の構造を壊すことなく、レンガを別の色や形のものに差し替えることができます。
• 透明性: 各レンガがどのように接続されているかを正確に確認できます。隠された層はありません。ロボットの学習方法を変更したい場合、特定の部品を一つ交換するだけで、残りはそのままです。

2. 「イン・バックワード」のトリック：歩きながら食べる

これらのロボットをトレーニングする際、最大の頭痛の種の一つはメモリです。部屋を横断しながら、その上にメモを書こうとする巨大な書類の束（勾配）を運ぶと想像してください。何かを行う前にその束を保持するには、非常に広いスペースが必要です。

torchtune は**「in-backward optimizer fusion（イン・バックワード・オプティマイザー・フュージョン）」**と呼ばれる巧妙なトリックを導入しています。

• 従来の方法: すべての書類を集め、それらを机まで運び、その後メモを書きます。これには巨大な机（メモリ）が必要です。
• torchtune の方法: 書類を拾う瞬間にその紙にメモを書き、すぐにその紙を捨てます。一度に束全体を保持する必要は決してありません。
• 結果: これにより、莫大な量のメモリが節約されます。本論文によれば、これはコンピューターがクラッシュする（メモリ不足）ことと、標準的なハードウェアで巨大なモデル（Llama 3.3 70B のようなもの）のトレーニングに成功することとの違いを生み出します。

3. 「ロス並列」のトリック：ケーキを切る

ロボットが自身のパフォーマンス（「ロス」）を計算すると、メモリを大量に消費する巨大で密集した数値の表を作成することがよくあります。

• アナロジー: 1,000 人分のケーキを一度に焼こうと想像してください。一つのオーブンには大きすぎます。
• 解決策: torchtune はケーキをより小さなスライスに切り分け、異なるオーブン（異なるプロセッサ）で同時に焼きます。巨大なケーキ全体を一つの場所に保持しようとは決してしません。これにより、システムはスペース不足に陥ることなく、巨大な語彙を持つモデルを処理できます。

4. 「非同期」工場：組立ライン

高度なトレーニング（強化学習など）において、ロボットは「考える」（回答を生成する）ことと「学ぶ」（脳を更新する）ことを行わなければなりません。通常、これらは工場のように、塗装ステーションがアセンブリラインが忙しい間、アイドル状態になるように、順番に起こります。

• torchtune のアプローチ: 彼らは非同期の組立ラインを構築しました。
• 仕組み: 作業員の一団が回答を生成する（塗装する）のに忙しくしている間、もう一団はすでにトレーニング（組み立て）に忙しく働いています。彼らはコンベアベルト（キュー）を使用して、作業を相互に渡します。これにより、工場全体が停止と開始を繰り返すのではなく、100% の能力で稼働し続けることができます。

5. 結果：速度と効率

著者らは、torchtune を他の人気のあるツール（Axolotl および Unsloth）と比較してテストしました。

• レース: 直接対決において、torchtune はトレーニングをより速く完了するか、より少ないメモリを使用することが多かったです。
• 「OOM（メモリ不足）」の修正: 最大のモデルにおいて、他のツールはメモリ不足により頻繁にクラッシュしました。torchtune は「歩きながら食べる」方法のようなメモリ節約のトリックを使用することで、他のツールが失敗した場所でこれらの巨大なモデルをトレーニングすることができました。
• 柔軟性: レゴのように構築されているため、研究者はこれらのトリックを組み合わせて使用できます。彼らは、すべてのトリックを一緒に使用することが最良の結果をもたらしましたが、必要に応じて単一のトリックのみを使用することも可能であることを発見しました。

まとめ

torchtune は、AI トレーニングを、ロックされたブラックボックスではなく、透明で交換可能なブロックのセットとして扱う新しいオープンソースツールキットです。データを保存するのではなく即座に処理することでメモリを節約し、タスクを並列実行することで速度を向上させ、プロセスのすべての部分を微調整するための研究者への完全な制御を提供します。この論文は、小規模な実験から大規模な産業規模のモデルトレーニングまで、既存のツールよりも優れていることを示しています。

技術概要

技術概要：torchtune – PyTorch ネイティブのポストトレーニングライブラリ

1. 問題定義

現代の大規模言語モデル（LLM）は、オープンウェイトモデルをダウンストリームタスクに適応させるために、教師あり微調整（SFT）、選好最適化、蒸留、および RL ベースのアライメントといった多段階のポストトレーニングパイプラインに大きく依存しています。しかし、この段階向けの既存フレームワークは、以下のような重大なトレードオフに直面しています：

• 複雑な依存関係スタック: transformers や関連ライブラリの上に構築されたフレームワークは、広範なトランジティブ依存関係を継承し、デプロイと再現性を複雑化しています。
• 密結合: モデル構築、トレーナーロジック、分散ポリシー、アダプター挿入は、しばしばファクトリレイヤー全体にわたって抽象化されており、基盤となる PyTorch モジュールを変更せずに微細な修正を行うことが困難です。
• 不均一なパフォーマンスアクセス: 汎用的な実装は、FSDP2、DTensor、torch.compile、ロス並列化などの現代の PyTorch パフォーマンスパスを活用できず、一方、カーネル特化型システムはトレーニングループの透明性を犠牲にすることが多いです。
• 断片化したサポート: 異なるポストトレーニングレシピ（SFT、DPO、PPO、GRPO、KD）はしばしば別々のライブラリに存在し、制御された比較を妨げています。
• 分散構成可能性: マルチノードトレーニング、テンソル並列化、コンテキスト並列化のサポートはフレームワーク間で一貫性が欠けており、スケールに応じて異なるバックエンドを必要とします。

2. 手法と設計原則

torchtune は、ポストトレーニングライフサイクルを合理化するために設計された PyTorch ネイティブのライブラリとして導入されます。モノリシックなトレーナーとは異なり、これは硬直した抽象化ではなく、構成可能なコンポーネントを中心に構築されています。

コアアーキテクチャ

• モジュール型コンポーネント: このライブラリは、モデルの組み立てをトレーニングロジックから分離します。モデルビルダーは明示的にトランスフォーマーブロックを構築するため、アーキテクチャのバリエーション（LoRA、量子化、カスタムアテンションカーネル）を、共有デコーダロジックやトレーニングレシピを書き換えることなく、局所的に差し替えることができます。
• YAML 駆動型レシピ: Hydra に着想を得て、レシピは YAML 設定でパラメータ化されたトレーニング手順（例：SFT、DPO、GRPO）を定義します。コンポーネント（モデル、データセット、オプティマイザ、ロス）は独立して差し替え可能です。コマンドライン上でのオーバーライドにより、スweep 型の実験が可能になります。
• ネイティブ PyTorch 実装: torchtune は、Llama や Qwen などの現代的なオープンソース LLM の純粋な PyTorch 参照実装を提供します。これらは transformers の対応する実装と数値的に同等ですが、読みやすく修正しやすいものです。これにより、Hugging Face Hub や TorchAO との相互運用性を維持しつつ、transformers のトレーニングループへの依存を排除します。

主要な技術的革新

1. バックワード内オプティマイザ融合:

   • メカニズム: 全パラメータの勾配を蓄積してから更新するのではなく、各パラメータの勾配が利用可能になった時点で即座にオプティマイザ更新を実行します。
   • 実装: 各パラメータごとに 1 つのオプティマイザオブジェクトをインスタンス化し、勾配蓄積後のフックを登録して、即座に step() と zero_grad() を呼び出します。
   • 利点: 勾配テンソルの生存期間を短縮し、ピーク勾配メモリを大幅に削減します。これは、限られたハードウェア上で大規模モデル（例：Llama 3.3 70B）を適合させる上で重要です。
   • 制約: 1 回のバックワードパスあたり 1 回のオプティマイザ更新（$K=1$）を前提としており、勾配蓄積が必要な場合はバッチサイズを調整する必要があります。

2. 線形クロスエントロピー（LCE）ロス:

   • メカニズム: 最終出力射影とクロスエントロピー計算を融合します。投影前に無視するトークンをマスクし、隠れ状態をチャンク単位で処理します。
   • 利点: 密な $[B, S, V]$ 対数尤度テンソルの具体化を防ぎ、特に大規模語彙においてロス計算中のピークメモリを削減します。PyTorch のロス並列コンテキストと構成可能です。

3. 構成可能な並列化スタック:

   • PyTorch のDTensor API に基づいて構築されています。
   • FSDP2（2D メッシュを備えたデータ並列化）、テンソル並列化、シーケンス並列化、およびエキスパート並列化（MoE 用）をサポートしています。
   • Ring Attention によるコンテキスト並列化を含みます。
   • ロス並列化は、完全な対数尤度の具体化を回避するために、語彙次元にわたって出力特徴をシャードします。

4. 非同期 GRPO:

   • 設計: Ray 調整キューとリプレイバッファを使用して、ロールアウト生成とポリシー更新を分離します。
   • アーキテクチャ: 推論（vLLM ベースのコレクター）、後処理（報酬計算）、トレーニング（分散ワーカー）を分離します。
   • モード: 同期交互、オンポリシー非同期オーバーラップ、および制限された遅延を伴う制御されたオフポリシーロールアウトをサポートします。

3. 実験結果

著者は、0.6B から 70B パラメータまでのモデル（Qwen3、Llama 3.3）を使用し、シングル GPU およびマルチ GPU（8x H100）設定において、torchtune をAxolotlおよびUnslothと比較評価しました。

主要な知見

• メモリ効率:
  • Optim Bwd: ベースライン構成では Out-Of-Memory（OOM）エラーが発生した環境で、8 個の H100 上での Llama 3.3 70B のトレーニングを可能にしました。
  • アクティベーションチェックポイント（AC）: 一貫してピークメモリを削減し、ベースラインでは失敗していた 8B モデルの実行を可能にしました。
  • 低ビットオプティマイザ: AdamW8Bit が最大の絶対メモリ削減を提供しました（例：Qwen3-1.7B は 11.7GB から 4.9GB に低下）。
  • 比較: 8B モデルでの DPO トレーニングにおいて、torchtune は標準的な AdamW を使用してメモリ内に収まりましたが、Axolotl は 8 ビットオプティマイザを必要としたか、完全に失敗しました。
• スループット:
  • コンパイル: torch.compile は、小規模から中規模のモデル（例：Qwen3-0.6B が 5.2k トークン/秒から 7.9k トークン/秒に増加）に対して最も信頼性の高いスループット向上を提供しました。
  • シーケンスパッキング: 有効なトークン利用率とスループットを大幅に向上させました（例：Qwen3-0.6B はパッキングにより 57k トークン/秒に到達）。
  • 相乗効果: 最適化は相補的であることが判明しました。コンパイルはスループットを駆動し、メモリ指向の技術（AC、Optim Bwd、LCE）はより大規模なスケールでの実現可能性を決定します。
• 柔軟性: このライブラリは、トレーニングループを書き換えることなく、フル微調整、LoRA、QLoRA、およびさまざまな並列化戦略を正常にサポートしました。

4. 意義と主張

本論文は、torchtuneを再現性のある LLM ポストトレーニング研究の実用的な基盤として位置付けています。その主な意義は以下の点にあります：

1. 透明性とハッカビリティ: 研究の表面を実行された PyTorch コードの近くに保つことで、研究者がトレーニングループを直接検査および修正することを可能にし、高レベルトレーナーの「ブラックボックス」性を回避します。
2. バランスの取れたトレードオフ: YAML レシピを通じた使いやすさ、ネイティブ PyTorch 最適化を通じたパフォーマンス、モジュール型コンポーネントを通じた拡張性のバランスに成功しています。
3. 統合フレームワーク: 多様なポストトレーニング手法（SFT、DPO、GRPO、KD）を単一の構成可能なスタックに統合し、異なるアルゴリズムや最適化戦略間の制御された比較を容易にします。

著者は、torchtune が迅速な実験と効率性の高いデプロイ指向ワークフローを可能にしつつ、迅速な研究反復に必要な柔軟性も維持しており、高レベルの自動化トレーナーと低レベルのパフォーマンス特化カーネルの間のギャップを効果的に埋めると主張しています。