PyTorch AutoSP：コンパイラベースのツールがトレーニングコードを自動的にシーケンス並列化、100k+トークンコンテキストに対応

2026年4月29日、PyTorchチームはUIUC SSAIL Lab、Anyscale、Snowflakeの研究者と共同でAutoSPを発表しました。これはDeepSpeed/DeepCompileエコシステム内のコンパイラベースのツールで、標準的な単一GPUトランスフォーマートレーニングコードをシーケンス並列形式に自動変換します。目標：分散コードを手動で実装することなく、超長コンテキスト（100k+トークン）のLLMトレーニングを可能にすることです。

AutoSPが解決する問題

長コンテキストのLLMトレーニングでは、トランスフォーマーの活性化がコンテキスト長に対して2乗スケーリングするため、複数のGPUにシーケンスを分散させる（シーケンス並列化、SP）必要があり、OOMエラーが簡単に発生します。既存のソリューション（RingFlashAttention、DeepSpeed-Ulysses）は、トークンパーティショニング、通信コレクティブ、複雑なアテンションマスキングなど、トレーニングコードの手動書き直しが必要です。

AutoSPはこれらすべてを自動化します：ユーザーが標準的なトランスフォーマーコードを書くと、コンパイラがSPを認識した形式に変換します。

有効化の方法

DeepSpeedの設定に3行追加するだけです：

config = {
    "compile": {
        "deepcompile": True,
        "passes": ["autosp"]
    },
    "sequence_parallel_size": 4,
    "zero_optimization": {"stage": 1},  # AutoSPはZeRO 0/1と組み合わせ可能
    ...
}

さらに入力準備のためのprepare_auto_sp_inputs()ユーティリティが必要です。内部戦略はDeepSpeed-Ulyssesアーキテクチャを採用し、NVLink/fat-treeネットワークでのGPU数増加に対して一定の通信オーバーヘッドを維持し、アテンションヘッド数（例：7-8Bモデルの32ヘッド）までスケールします。

シーケンス認識アクティベーションチェックポイント（SAC）

AutoSPはSAC（長コンテキストトレーニング向けに最適化されたカスタムチェックポイント戦略）も導入しています。保守的なPyTorch 2.0の最大フロー最小カット定式化とは異なり、SACは長コンテキストの特有のFLOPダイナミクスを活用し、計算コストの低い演算子からの中間活性化を解放し、バックワードパスで再計算します。トレードオフ： わずかにスループットが低下しますが、さらに長いコンテキストを実現可能にします。

結果

8×A100-80GB SXMノード（PyTorch 2.7、CUDA 12.8）でLlama 3.1の各サイズモデルをテスト：

• 同じリソースで最大トレーニング可能なシーケンス長が大幅に向上
• 手書きのRingFlashAttentionとDeepSpeed-Ulyssesベースラインと比較してランタイムオーバーヘッドは最小限

エンドツーエンドの例（Llama 3.1 8Bを含む）はgithub.com/deepspeedai/DeepSpeedExamples/tree/master/benchmarks/autospで利用可能です。

制限事項

現在AutoSPは単一のコンパイル可能なアーティファクト（トランスフォーマー全体を1つのブロックとしてコンパイル）を必要とし、モデル内でのグラフブレークを許可しません。チームはグラフブレーク耐性の拡張を次のステップとして挙げています。