Ulyssesシーケンス並列

Ulysses sequence parallelism (SP) trains on very long sequences by splitting them across multiple GPUs. To compute attention correctly, an all-to-all collective swaps the sharding dimension from sequence to attention heads. Each GPU then has the full sequence and computes attention locally over a subset of heads. A second all-to-all returns to the sequence-sharded layout so the rest of the forward pass continues locally on each chunk.

                        GPU 0                       GPU 1
                   ┌───────────────┐          ┌───────────────┐
  forward          │ tokens 0..N/2 │          │ tokens N/2..N │  ← each GPU holds half the sequence
  (seq-sharded)    │  all H heads  │          │  all H heads  │
                   └───────┬───────┘          └───────┬───────┘
                           └───────── all-to-all ──────┘
                   ┌───────────────┐          ┌───────────────┐
  attention        │ all N tokens  │          │ all N tokens  │  ← now each GPU has the full sequence
  (head-sharded)   │ heads 0..H/2  │          │ heads H/2..H  │  ← but only half the heads
                   └───────┬───────┘          └───────┬───────┘
                           └───────── all-to-all ──────┘
                   ┌───────────────┐          ┌───────────────┐
  forward          │ tokens 0..N/2 │          │ tokens N/2..N │  ← back to seq-sharded
  (seq-sharded)    │  all H heads  │          │  all H heads  │
                   └───────────────┘          └───────────────┘

このガイドでは、ALST（Arctic Long Sequence Training）のUlyssesシーケンス並列化コンポーネントについて説明します。ALSTシステム全体には、Transformersでは利用できないTiledMLPとアクティベーションチェックポイントのオフローディングも含まれています。完全なシステムについては、DeepSpeed ALSTチュートリアルを参照してください。

設定

シーケンス並列化には、Accelerate v1.12.0 以上と 2 枚以上の GPU が必要です。Accelerate の ParallelismConfig でシーケンス並列化を設定し、parallelism_config または Accelerate 設定ファイル に渡します。

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_replicate_size=1,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",
    ),
)

training_args = TrainingArguments(
    ...,
    deepspeed="path/to/deepspeed_config.json",
    parallelism_config=parallelism_config,
)

accelerate launch --num_processes 4 train.py \
--output_dir output_dir \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 1

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 4
parallelism_config:
  parallelism_config_sp_size: 4
  parallelism_config_dp_replicate_size: 1
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: flash_attention_2

accelerate launch --config_file alst_config.yaml train.py \
--output_dir output_dir \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 1

シーケンス並列性を設定するには、以下のフィールドが重要です。

Trainer は、DataLoader のシャーディング、position_ids の生成、ラベルシフト、および SP ランク間の損失集約を自動的に処理します。カスタムの学習ループを実装する場合は、代わりに Accelerate の Sequence Parallelism ガイドを参照してください。

• Ulyssesシーケンス並列を使用するには、sp_backendを"deepspeed"に設定する必要があります。

• sp_size は、単一のシーケンスを並列で処理する GPU の数です。テンソル並列ではすべてのランクが同一のデータを受信するのに対し、各 SP ランクは DataLoader から固有のデータストリームを受信します。実効的な dp_world_size = world_size / sp_size となるため、4 つの GPU で sp_size=4 の場合、バッチサイズ計算用の dp_world_size=1 になります。また、シーケンスは sp_size の倍数にパディングする必要があります。データコラテーターで pad_to_multiple_of を適切に設定してください。

  アテンションヘッドの数はsp_sizeで割り切れる必要があります。32個のヘッドを持つモデルでは、sp_sizeは1、2、4、8、16、32のいずれかをサポートします。

  コピーしました

  from transformers import DataCollatorForLanguageModeling
  
  data_collator = DataCollatorForLanguageModeling(
      tokenizer=tokenizer,
      mlm=False,
      pad_to_multiple_of=sp_size,
  )

• sp_seq_length_is_variable は、可変シーケンス長の処理を制御します。バッチ間で長さが異なる場合は True に設定してください(推奨)。すべてのシーケンスが sp_seq_length で指定された固定長にパディングされる場合は、False に設定してください。

• sp_attn_implementation はアテンションのバックエンドを設定します。サポートされている値は "sdpa"、"flash_attention_2"、または "flash_attention_3" です。FlashAttention が推奨されます。特に、バッチ内で複数のサンプルをパッキングする場合です。サンプルをパッキングしている場合、SDPA はサンプルの境界をまたいで誤ってアテンションを行う可能性があります。Eager attention はサポートされていません。これは、メモリとスケーリングの理由から、4 次元の attention_mask が破棄されるためです。

データ並列との組み合わせ

シーケンス並列化とデータ並列化は同じGPUを使用し、SPには追加のハードウェアは不要です。両方を同時に実行するには、dp_replicate_sizeまたはdp_shard_sizeを設定し、dp_replicate_size × dp_shard_size × sp_sizeがGPUの総数に等しくなるようにしてください。

例えば、8つのGPUとsp_size=4を使用する場合、dp_replicate_size=2を設定します（2 × 1 × 4 = 8）。

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_replicate_size=2,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",
    ),
)

次のステップ

• FSDP2 によるコンテキスト並列化（cp_size）は、Ulysses（sp_size）とは独立しています。コンテキスト並列化での評価を参照してください。
• Accelerate の シーケンス並列化ガイドでは、Ulysses の実装についてより詳細に説明し、カスタムな学習ループの書き方も示しています。
• DeepSpeed ALST チュートリアルでは、TiledMLPやアクティベーションチェックポイントのオフロードを含む、ALSTシステムの全体を扱っています。
• 並列化手法のガイドでは、シーケンス並列を ZeRO などの他の戦略と組み合わせる方法について説明しています。
• Ulysses Sequence Parallelism: Training with Million-Token Contexts というブログ記事では、Ulysses の仕組みと、Accelerate、Trainer、SFTTrainer への統合方法について説明しています。