AutoSP の紹介：長文コンテキスト LLM 学習を自動化する新技術

近年、大規模言語モデル（LLM）は、トークン数が 10 万を超えるような極めて長いコンテキストを扱うタスク向けに訓練されるようになっています。このようなトークン数では、ZeRO や FSDP といった従来の訓練手法を用いてデバイス数をスケールさせても、メモリ不足（OOM: out-of-memory）の問題が表面化し始めます。これらの問題を回避するために、シーケンス並列処理（SP: sequence parallelism）という手法が広く用いられています。これは、入力トークンを複数のデバイスに分割することで、GPU 数の増加に伴って長いコンテキストの訓練を可能にする並列訓練技術です。 しかし、SP の実装は非常に困難で知られており、DeepSpeed や HuggingFace といった既存ライブラリに対して侵入的なコード変更が必要です。これらのコード変更には、入力トークンのコンテキスト（および中間活性化）の分割、通信集合（collectives）の挿入、計算と通信の重なり付けが含まれ、これらは順伝播（forward pass）と逆伝播（backwards pass）の両方で行わなければなりません。その結果、長いコンテキスト機能を試したい研究者たちは、そのような機能を実現するためにシステムスタックのエンジニアリングに多大な労力を費やすことになり、異なるハードウェアベンダーごとにこの作業を繰り返さねばなりません。 この複雑さを回避するため、[AutoSP](https://openreview.net/pdf?id=0fgsHvmBBI) を導入します。これは完全自動化されたコンパイラベースのソリューションであり、記述が容易なトレーニングコードを、より長い入力コンテキストで効率的に GPU を使用してトレーニングしつつ、既存のパラレル戦略（ZeRO など）と組み合わせ可能なマルチ GPU シーケンスパラレルコードに変換するものです。これにより、開発者が長文コンテキストトレーニングのためにトレーニングパイプラインを繰り返し修正するという面倒な必要性が不要になります。ユーザーは now AutoSP をインポートし、AutoSP バックエンドを使用して任意のモデルをコンパイルするだけでよく、誰でも長文コンテキストトレーニングのパワーを利用できるようになります。さらに、この技術をコンパイラに埋め込むことで、当社のアプローチはパフォーマンスポータブルとなります：高性能なシーケンスパラレル（Sequence Parallelism）が多様なハードウェア上で実現可能です。 本記事の構成は以下の通りです。(1) AutoSP とモデル科学者が長文コンテキストトレーニングを可能にするための使用方法、(2) AutoSP の主要な設計判断、(3) 使いやすさと影響を示す主要な AutoSP の結果、(4) いくつかの制限事項と AutoSP が実行できないこと。 ## AutoSP の使い方 AutoSP の重要な設計哲学は、複数の GPU をプログラミングする際の複雑さのほとんどをユーザーから抽象化し、シンプルさを保つことにあります。これを実現するために、AutoSP は [DeepCompile](https://arxiv.org/pdf/2504.09983) 内に実装されています。これは DeepSpeed 内のコンパイラエコシステムであり、深層ニューラルネットワークのトレーニングに対して多様な最適化をプログラム可能にします。これにより、DeepSpeed を使用するあらゆるユーザーが、ほとんど手間をかけずにシーケンス並列処理（Sequence Parallelism）を自動的に有効にできます。次に例を見てみましょう。 # DeepSpeed の設定をインスタンス化します。 # 8 つの GPU を使用し、2 つのデータ並列（DP）ランクと 4 つのシーケンス並列（SP）ランクがあると仮定します。 config = { "train_micro_batch_size_per_gpu": 1, "train_batch_size": 2, "steps_per_print": 1, "optimiser": { "type": "Adam", "params": { "lr": 1e-4 } }, "zero_optimization": { "stage": 1, # AutoSP は ZeRO 0/1 と相互運用します。 }, # DeepCompile をオンにし、AutoSP パスをトリガーするように設定するだけです。 "compile": { "deepcompile": True, "passes": ["autosp"] }, "sequence_parallel_size": 4, "gradient_clipping": 1.0, } # モデルで DeepSpeed を初期化します。 model, _, _ = deepspeed.initialize(config=config,model=model) # モデルをコンパイルし、自動的に AutoSP パスを適用します。 model.compile(compile_kwargs={"dynamic": True}) for idx, batch in enumerate(train_loader): # Custom function that we expose within: # deepspeed/compile/passes/sp_compile. inputs, labels, positions, mask = prepare_auto_sp_inputs(batch) loss = model( input_ids=inputs, labels=labels, position_ids=positions, attention_mask=mask ) ... # Backwards pass, optimiser step etc... 上記の例から分かるように、ユーザーは単一デバイス上で動作する既存のトレーニングコードを引き続き使用し、以下の手順を実行します。(1) DeepSpeed で公開されている `prepare_autosp_input` ユーティリティ関数を使用して、AutoSP 内のプログラム解析で利用するための入力トークン、アテンションマスク、位置 ID に軽量なタグ付けを行います。(2) DeepSpeed の設定を調整して DeepCompile を有効にし、「passes」フラグを「autosp」に指定します。残りの処理は、モデルのコンパイル時に呼び出される AutoSP コンパイラパスによって自動的に処理され、他の長文脈トレーニング最適化と並列してシーケンス並列化が自動的に有効化されます。AutoSP はさらに ZeRO stage 1 ともデフォルトで自動連携し、DeepSpeed 設定において AutoSP フラグと共に ZeRO-1 フラグを設定するだけで、両方の戦略を組み合わせることができます。 ## AutoSP コンパイラパス AutoSP はユーザーコードを変換してより長い文脈でのトレーニングを可能にするため、ここでは AutoSP の主要な設計ポイントとコード変換の内容、および透明性を保つためのユーザーへの影響について簡単に解説します。 シーケンシャル並列化コード変換。AutoSP は自動的にシングル GPU コードをマルチ GPU 用のシーケンシャル並列化（Sequence Parallelism、以下 SP）コードに変換します。AutoSP がコードを変換する具体的な SP ストラテジーは [DeepSpeed-Ulysses](https://dl.acm.org/doi/10.1145/3662158.3662806) です。他のストラテジー（例：[RingAttention](https://arxiv.org/pdf/2310.01889)）と比較して DeepSpeed-Ulysses に焦点を当てているのは、NVLink ネットワークトポロジやファットツリーネットワークにおいて、GPU 数が増加しても通信オーバーヘッドが一定に保たれるためです。ただし、DeepSpeed-Ulysses はモデル内のヘッド数（7-8B モデルでは 32）まで SP サイズを拡張できるという点に限界があります。 より長いコンテキストでのトレーニングのための活性化チェックポイント。AutoSP は、長文コンテキストモデリング向けに特別に設計されたカスタム活性化チェックポイント (AC) 戦略も適用しています。AC は計算コストの低い演算子の中間活性化を解放し、必要な勾配を計算するために逆伝播パスで必要に応じて再計算を行います。PyTorch-2.0 では、自動化された最大フロー最小カットに基づく [AC の定式化](https://dev-discuss.pytorch.org/t/min-cut-optimal-recomputation-i-e-activation-checkpointing-with-aotautograd/467) が導入されていますが、長文コンテキストモデリングにおいては過度に保守的であることが判明しました。そこで我々は、長文コンテキストトレーニングを対象とした新たな AC 戦略として、独特な長文コンテキストの FLOP(浮動小数点演算数) ダイナミクスを活用する「シーケンス認識型 AC (SAC)」を導入しました。(AutoSP のデフォルト設定で) トリガーされると、これはトレーニングのスループットをわずかに低下させます。しかし、これなしでは長いコンテキストでのトレーニングは不可能となるため、ユーザーは OOM(メモリ不足) が発生する構成に対してのみ、このパスを選択的にオンにすることができます。 ## 実モデルにおける AutoSP の評価 AutoSP の実用性を示すために、異なるサイズのモデルに対して NVIDIA GPU で性能を評価し、使いやすさがランタイムパフォーマンスにほとんどあるいは全くコストをかけないことを示します。8 個の A100-80Gb SXM ノード上で、さまざまな Llama 3.1 モデルをベンチマークしました。PyTorch 2.7 と CUDA 12.8 を使用し、AutoSP を以下の torch コンパイル済み手書きベースラインと比較します：[RingFlashAttention](https://github.com/zhuzilin/ring-flash-attention)、DeepSpeed-Ulysses、および ZeRO-3。主要な結果を以下にまとめます。 ![image](https://pytorch.org/wp-content/uploads/2026/04/unnamed-32.png) AutoSP は、同じリソース条件下で最大学習可能シーケンス長を増加させるだけでなく（左図 – 高いほど良い）、これらの利点はランタイムパフォーマンスにほとんどコストをかけないことも示されています（右図 – 低いほど良い）。 ## 制限事項 AutoSP には2つの主要な制限があります。第一に、ユーザーはトランスフォーマーを単一のコンパイル可能なアーティファクトとして強制的にコンパイルする必要があります。場合によっては、PyTorch ユーザーが多数の関数を個別にコンパイルしてそれらを1つのモデルに結合することがありますが、これは AutoSP では許可されていません。入力シーケンスを正しくシャードし、この情報をグラフ全体に伝播させるためには、モデル全体をコンパイルして確認する必要があるためです。第二に、コンパイル可能なアーティファクト内のグラフブレイクはすべて禁止されています。これにより情報の分析と伝播が複雑化するため、AutoSP をグラフブレイク耐性のあるものへと拡張することは今後の研究課題として残します。 ## 結論 AutoSP を使用すると、ユーザーは任意のトランスフォーマー学習コードにシーケンス並列化を容易に追加し、長期コンテキスト学習を強化するためのカスタム AC 戦略を実装できます。DeepSpeed との統合により、設定ファイルを変更するだけで、既存の DeepSpeed 学習コードを使用してより長いコンテキストで学習することが可能になります。実際のモデルワークロード（例：Llama 3.1 8B）で試せるエンドツーエンドのサンプルを [こちら](https://github.com/deepspeedai/DeepSpeedExamples/tree/master/benchmarks/autosp) に用意しています。長期コンテキスト学習がいかに容易になったか、ぜひお試しください。

AutoSP の紹介：長文コンテキスト LLM 学習を自動化する新技術

背景や根拠まで確認しますか？

調べる

選ぶ

サイト

背景や根拠まで確認しますか？

ニュースの次に確認する

調べる

選ぶ

サイト