Ulysses シーケンス並列化による百万トークンコンテキストのトレーニング

Hugging Face は、Snowflake AI Research の Arctic Long Sequence Training プロトコルの一部である Ulysses Sequence Parallelism を Accelerate、Transformers、TRL ライブラリに統合しました。この統合により、開発者は注意計算を複数の GPU に分散させ、注意機構の二次的なメモリスケーリングを克服することで、数十万トークンから数百万トークン規模のシーケンス上で大規模言語モデル（LLM）をトレーニングできるようになります。

長シーケンストレーニングの課題

標準的な Transformer の注意は、シーケンス長 $n$ に対して FLOPs とメモリの両方が二乗 ($O(n^2)$) でスケールします。FlashAttention はメモリ使用量を $O(n)$ に削減しますが、$O(n^2)$ の計算量は依然として残ります。シーケンスが 32k トークンを超える場合、トレーニングは通常単一 GPU のメモリ容量を超えてしまい、データ並列だけに頼らずシーケンス自体を複数デバイスに分割する手法が必要になります。

Ulysses シーケンス並列化の仕組み

Ulysses Sequence Parallelism（SP）は、シーケンス次元と注意ヘッドの両方を GPU に跨って分割することで注意計算を分散します。プロセスは以下の手順で進みます。

  1. シーケンスシャーディング: 入力シーケンスを $P$ 個の GPU に分割し、各 GPU がトークンのローカルチャンクを保持します。
  2. QKV 投影: 各 GPU がローカルチャンクに対して query、key、value の投影を計算します。
  3. All-to-All 通信: All-to-All の集合操作でデータを再配布し、各 GPU がすべてのシーケンス位置を保持しますが、注意ヘッドはサブセットに限定されます。
  4. ローカル注意: GPU は FlashAttention または SDPA を用いて割り当てられたヘッドの注意を計算します。
  5. All-to-All 通信: 2 回目の All-to-All 操作でデータをシーケンスシャーディング形式に戻します。
  6. 出力投影: 各 GPU がローカルシーケンスチャンクに対して出力投影を計算します。

通信複雑度

Ulysses は注意レイヤーごとに 2 回の All-to-All 操作が必要で、GPU あたりの通信量は $O(n \cdot d / P)$（ここで $n$ はシーケンス長、$d$ は隠れ次元、$P$ は並列度）です。これは Ring Attention の $O(n \cdot d)$（GPU あたり）で、$P-1$ ホップにわたって転送を直列化する方式よりも効率的です。

エコシステムへの統合

Accelerate

Accelerate は ParallelismConfig クラスと DeepSpeed 統合を通じて Ulysses を実装します。主なパラメータは sp_size（シーケンス並列に使用する GPU 数）と sp_backend（"deepspeed" に設定する必要があります）です。accelerator.prepare() が呼び出されると、システムは自動的にモデルを UlyssesSPAttentionHF に登録し、データローダーを UlyssesSPDataLoaderAdapter でラップします。

Transformers Trainer

Transformers の Trainer は TrainingArguments.parallelism_config を介して Ulysses 統合を処理します。データローダーのラップ、シーケンスシャーディング、加重ロス集約を自動化し、トークンがランク間で不均等に分配されても勾配が正しく計算されるようにします。

TRL SFTTrainer

TRL の SFTTrainer は、パディング無駄を削減する packing 機能など、教師ありファインチューニング向けの最適化を追加します。pad_to_multiple_of は sp_size と同じに設定する必要があり、シーケンスの除算可能性を保証します。Ulysses が有効な場合、SFTTrainer は事前シフトされたラベルも自動的に管理します。

Ulysses と Ring Attention の比較

項目 Ulysses (DeepSpeed) Ring Attention (FSDP2)
並列化手法 注意ヘッドの分割 リングベースの KV 交換
バックエンド DeepSpeed ZeRO PyTorch FSDP2
注意サポート FlashAttention 2/3、SDPA SDPA のみ
通信方式 レイヤーごとに 2 回の all-to-all P2P リング通信
GPU あたりの通信量 $O(\text{total_seq} \times \text{hidden} / \text{sp_size})$ $O(\text{total_seq} \times \text{hidden})$
ヘッド数制約 num_heads >= sp_size なし

パフォーマンスベンチマーク

Hugging Face は、Qwen3-4B を Gutenberg 英語データセット上で H100 80GB GPU を使用して Ulysses SP のベンチマークを実施しました。

メモリ削減

SP=4 を使用すると、同一シーケンス長で GPU あたりのメモリ使用量が 3.3 倍削減されます。これにより、ベースラインの 8K トークン（DP=4）から 96K トークン（SP=4）へスケールアップでき、80GB のメモリ上限内に収まります。128K トークンでは OOM（Out‑of‑Memory）状態に達しました。

スループット

スループットはシーケンス長が増加するにつれて向上します。二次的な注意計算が通信オーバーヘッドを支配するためです。64K トークンで SP=4 は 13,396 トークン/秒を達成し、8K ベースラインの 3.7 倍のスループットとなります。

実装のベストプラクティス

  • シーケンス除算可能性: pad_to_multiple_of を使用してシーケンス長が sp_size で割り切れるようにします。
  • 注意バックエンド: Ampere GPU では FlashAttention 2、Hopper GPU では FlashAttention 3 を使用します。
  • メモリ最適化: Ulysses と DeepSpeed ZeRO Stage 3 を組み合わせ、環境変数 PYTORCH_ALLOC_CONF=expandable_segments:True を設定して断片化を削減します。
  • 2D 並列化: GPU 数に応じて sp_size と dp_shard_size をバランスさせ、最大シーケンス長または高スループットのどちらかを最適化します。
  • 追加カーネル: Liger‑Kernel の FusedLinearCrossEntropy と TiledMLP を使用して、ロス計算や大規模行列演算時の作業メモリをさらに削減します。

Sources