LLMを10万トークン超の長文で学習させようとすると、メモリ不足で止まる。この問題を設定ファイルの変更だけで解決するコンパイラベースのツールが登場しました。

この記事でわかること

  • AutoSPが解決する長文コンテキスト学習の課題
  • 従来のシーケンス並列化との違い
  • 導入手順と必要なコード変更
  • NVIDIA・AMDでの性能検証結果

長文コンテキスト学習の壁

Introducing AutoSP – PyTorch
pytorch.org

2026年4月29日、PyTorch公式ブログでAutoSPが発表されました。イリノイ大学アーバナ・シャンペーン校（UIUC）のSSAIL Lab、Anyscale、Snowflakeの研究チームが開発し、ICLR 2026で採択された研究成果です。

LLMの学習では、入力トークン数が10万を超えると、ZeROやFSDPといった既存の分散学習手法でもGPUメモリが足りなくなります。この問題に対応する手法がシーケンス並列化（Sequence Parallelism、以下SP）です。入力トークンを複数のGPUに分割して処理します。

しかし、SPの実装はこれまで非常に手間がかかるものでした。DeepSpeedやHugging Faceのコードに手を入れて、トークンの分割、通信処理の挿入、計算と通信の重複処理を前方・後方パスの両方で実装する必要があります。ハードウェアが変われば、この作業をやり直すことになります。

AutoSPの仕組み

AutoSPはコンパイラベースのアプローチでこの問題を解決します。単一GPUで動く通常の学習コードを、自動的にマルチGPUのSPコードに変換します。

内部ではDeepSpeedの「DeepCompile」というコンパイラ基盤の上に構築されています。DeepCompileは単一GPUのモデル実装を中間表現（IR）に変換し、最適化パスを通じて分散学習に必要な操作を自動的に挿入する仕組みです。AutoSPはこのパスの一つとして動作します。

AutoSPが適用するSP戦略はDeepSpeed-Ulyssesです。各GPUがシーケンスの一部を受け持ち、Attention計算の直前にall-to-all通信でクエリ・キー・バリューを再配置します。各GPUは全シーケンスの一部のAttentionヘッドを担当する形です。RingAttentionと異なり、NVLinkやファットツリーネットワーク上でGPU数を増やしても通信オーバーヘッドが一定に保たれる点が特徴です。

さらにAutoSPは、長文学習に特化した活性化チェックポイント戦略「SAC（Sequence-aware AC）」も自動で適用します。PyTorch 2.0のmax-flow min-cutベースのACは長文学習には保守的すぎるため、長文特有のFLOP特性を活用した独自の戦略を採用しています。スループットへの影響はわずかで、これがないとそもそも長文での学習ができないケースに対応します。

導入に必要な変更

AutoSPの導入に必要な変更は2点だけです。

まず、DeepSpeedの設定ファイルにDeepCompileとAutoSPのフラグを追加します。compileセクションでdeepcompileをオンにし、passesにautospを指定します。sequence_parallel_sizeでSPに使うGPU数を設定します。ZeROステージ1との併用も、設定ファイルの変更だけで対応します。

次に、入力データの前処理でprepare_auto_sp_inputsというユーティリティ関数を呼び出します。トークン、Attentionマスク、位置IDに軽量なタグ付けを行うだけで、残りはコンパイラパスが自動で処理します。

モデルのコード自体を書き換える必要はありません。既存のDeepSpeed学習コードがそのまま使えます。

性能検証の結果

研究チームはLlama 3.1モデルを使い、A100-80GB SXM 8台のノードで検証しています。PyTorch 2.7、CUDA 12.8の環境で、手書きのRingFlashAttention、DeepSpeed-Ulysses、ZeRO-3と比較しました。

AutoSPはNVIDIA GPU上で最大学習コンテキスト長を最大2.7倍に伸ばしました。AMD GPU上でも2.5倍の改善を達成しています。ランタイム性能への影響はほぼありません。コンパイラベースのため、ハードウェアを問わず高い性能を発揮できる点がポイントです。

現時点の制約

AutoSPには2つの制約があります。

1つ目は、Transformerモデル全体を単一のコンパイル対象として渡す必要がある点です。複数の関数を個別にコンパイルして組み合わせるパターンには対応していません。AutoSPはモデル全体を見てシーケンスの分割情報を伝播させるため、全体が一つのグラフである必要があります。

2つ目は、コンパイル対象にグラフブレイク（コンパイルの中断点）を含めない点です。グラフブレイクがあると情報の伝播が複雑になるため、現時点では非対応です。

まとめ

AutoSPは、長文コンテキストでのLLM学習を「設定ファイルの変更だけ」で実現するコンパイラ技術です。DeepSpeedのエコシステムに統合されており、既存の学習コードへの影響を最小限に抑えながら、学習可能なコンテキスト長を大幅に伸ばせます。Llama 3.1 8Bを使ったエンドツーエンドのサンプルコードがDeepSpeedのリポジトリで公開されています。長文対応の学習環境を構築したい開発者にとって、まず試す価値のある選択肢です。