🤗 accelerate におけるシーケンス並列

このガイドでは、🤗accelerateでシーケンス並列性を使用する基本的な内容について説明します。

関連性の非常に高いContext Parallelismも参照してください。

なぜシーケンス並列性が必要なのか

大規模言語モデル、そして最近の推論モデルの登場により、シーケンス長は急速に伸びています。これにアテンションの二次的なメモリ複雑度が加わったことで、長いシーケンスを持つモデルを訓練するためのより効率的な手法の必要性が高まっています。 シーケンス長が128kの場合、バニラなアテンション実装を前提とすると、bf16精度におけるアテンション行列のメモリ要件は128k * 128k * 2 bytes * num_heads = ~32 GB * num_headsとなります。確かに、これらのアテンション重みを物化しないflash attentionやSDPAを使用すれば、メモリ使用量は大幅に減少しますが、メモリ要件の増加は依然として無視できない規模です。

Ulysses Sequence parallelismにより、attention計算の入力をシーケンス次元に沿ってシャード化し、attentionを通常通り計算しながら、各GPUではattentionヘッドの一部のみを使用できます。これにより、いくつかの追加ツールを用いることで、15M以上のシーケンス長にスケールしてモデルをトレーニングできます。Ulysses SPにTiledMLP、Liger-Kernel、CPUへのActivation checkpoint offload、およびその他のいくつかのテクニックを組み合わせる方法については、論文Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequencesを参照してください。

Ulysses SPとFSDP CPの違いは何ですか

ドキュメントContext Parallelismでは、シーケンス次元でスライスする点はこの技術と同様ですが、ヘッド次元でのスライスの代わりにRing Attentionを使用する別の技術であるContext Parallelismのデプロイ方法について学ぶことができます。

以下の記事では、2つの技術の違いについて非常に詳細に説明しています。

• https://insujang.github.io/2024-01-11/tensor-parallelism-and-sequence-parallelism-detailed-analysis/
• https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention

ある記事から引用した簡単なまとめ:

• Ulysses SPは通信オーバーヘッドが比較的低い一方で、Attention Headsの数に制約があるため、ネットワークトポロジーに対して特定の要件（単一のレプリカについて、attention headsの数が参加するGPUの数で割り切れること）があります。All-to-all通信はレイテンシに敏感であり、Deepspeedが必要です。
• FSDP CP Ring-AttentionのP2Pリング通信には、前述の割り切れる要件はありませんが、通信量が多くなります。

最後に、論文USP: A Unified Sequence Parallelism Approach for Long Context Generative AIで説明されているように、SP + CPを組み合わせることで、さらに長いシーケンス長をサポートできるようになるはずです。ただし、これはまだ🤗accelerateに統合されていません。

サポートされているシーケンス並列化バックエンド

現在、シーケンス並列化のバックエンドとして利用可能なものはdeepspeedのみです。これは、Arctic Long Sequence Training technologyの一部であるモダン化された Ulysses SP に由来しています。ご自身のコードに直接統合したい場合は、チュートリアルも用意されています。

シーケンス並列性を使用するにはどうすればよいですか？

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

+# Example: 4 GPUs with sp_size=4, dp_shard_size=1
+# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
+     sp_backend="deepspeed",
+     sp_size=4,
+     dp_shard_size=1,  # Explicit: no data parallelism
+     sp_handler=DeepSpeedSequenceParallelConfig(
+         sp_seq_length_is_variable: true,
+         sp_attn_implementation="sdpa",
+     ),
+ )

accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

🤗accelerateの他の機能と同様に、accelerate launchに該当するフラグを渡すことでシーケンス並列化を有効にできます。この場合も変わりはありません:

accelerate launch --parallelism-config-sp-size 8  ...

sp_size やその他の設定を accelerate config コマンドで設定することもできます。これにより、設定が accelerate の設定ファイルに保存されるため、スクリプトを起動するたびに毎回引数を渡す必要がなくなります。

sequence parallelism は data parallelism と組み合わせられます。追加の GPU は不要です。 したがって、8 つの GPU を持っている場合、--parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8 と指定できます。または、ParallelismConfig クラスを使用してプログラム的に設定することもできます。

重要: dp_replicate_size × dp_shard_size × sp_size = num_processes であることを確認する必要があります。例えば、8つのGPUでsp_size=8の場合、dp_shard_size=1が必要です（1 × 1 × 8 = 8のため）。4つのGPUでsp_size=2の場合、2D並列化のためにdp_shard_size=2を使用できます（1 × 2 × 2 = 4のため）。

ALST/Ulysses SPバックエンドの構成

ALST/UlyssesSP は、この論文で説明されているように、アテンションヘッド並列性を用いてシーケンス並列性を実装しています。簡便さのため、シーケンス並列性の概念とセットアップを再利用しますが、ユーザーの視点からは同じです: 複数の GPU が 1 つのバッチの処理に使用されます。

ALSTが何を実現したかを示すために、単一のH100 GPUで500Kトークン、単一ノードで3.7Mトークン、そしてLlama-8Bではわずか4ノードで15Mトークンというシーケンス長でbf16での学習が可能になりました。HF Accelerateのこの機能は、ALSTの3つのコンポーネントのうち1つのみを有効にするため、達成可能なシーケンス長は小さくなります。ALSTの全機能を発揮するには、TiledMLP、CPUへのActivation checkpoint offload、およびその他のいくつかの機能の有効化が必要です。詳細については、このチュートリアルを参照してください。

deepspeedバックエンドを設定するには:

# Example: 4 GPUs with sp_size=4, dp_shard_size=1
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_shard_size=1,  # Explicit: no data parallelism
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=256,
        sp_seq_length_is_variable=True,
        sp_attn_implementation="sdpa",
    ),
)
accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

• sp_backend: ここではdeepspeedに設定します
• sp_sizeはシーケンス並列の度合いです。上記の例では4であるため、単一のバッチを処理するために4つのGPUが使用されます（同じGPU上でDP=4のデータ並列も実行されます）。
• sp_seq_length と sp_seq_length_is_variable はシーケンス長を扱うために使用されます。sp_seq_length_is_variable=True の場合、バックエンドはバッチ間で変化し得るシーケンス長で動作し、この場合 sp_seq_length の値はシーケンス並列度で割り切れる値に設定するか、全く設定しなくても構いません。この場合、各 forward においてシーケンス変数は入力から導出されます。False の場合、seq_length はバッチのシーケンス長次元と一致する必要があり、その次元は常に同じになるようパディングする必要があります。デフォルトは True です。
• sp_attn_implementation は sdpa、flash_attention_2、または flash_attention_3 のいずれかです。このシーケンス並列実装は attention_mask の代わりに position_ids を使用するため、eager は position_ids での動作をサポートするまでここでは機能しません。また、sdpa は複数のサンプルを1つに結合した場合を正しく処理しない点にご注意ください。結合されたサンプル全体を1つとしてアテンションを行います。サンプルが結合されていない場合は、sdpa は正しく動作します。したがって、常に動作する Flash Attention が理想的な選択です。

これらの値をDeepSpeedSequenceParallelConfigオブジェクトに設定する代わりに、環境変数を使用して同じことを実現することもできます。上記のリストの末尾に対応するものが以下の通りです。

• PARALLELISM_CONFIG_SP_BACKEND
• PARALLELISM_CONFIG_SP_SEQ_LENGTH
• PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE
• PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION

コードで渡されない場合、sp_size は --parallelism_config_sp_size という CLI 引数で設定できます。他の引数も同様です。また、accelerate の設定ファイル形式で設定することもできます。例えば、2 つの GPU の場合:

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 2
parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1  # Must satisfy: 1 × 1 × 2 = 2 num_processes
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: sdpa

前述したように、Ulyssesシーケンス並列化は通常、データ並列化と重ねて適用されます。つまり、同じランクがユニークなデータストリームを供給する役割と、Ulyssesシーケンス並列化の実行を兼ねます。しかし、次のようにレプリカを作成することもできます。

# Example: 4 GPUs with 2D parallelism (SP=2, DP=2)
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 2 × 1 × 2 = 4 GPUs
parallelism_config = ParallelismConfig(
    dp_replicate_size=2,
    dp_shard_size=1,  # Explicit: no sharding within replicas
    sp_size=2,
    sp_backend="deepspeed",
    sp_handler=DeepSpeedSequenceParallelConfig(...),
)

ここでは4つのGPUを使用し、シーケンス並列性のレプリカを2つ設定しています。データ並列性を駆動するのはDeepspeed-ZeROです。

UlyssesSPDataLoaderAdapter の内部には多くの処理が隠されています。これは裏側で動作し、元の DataLoader オブジェクトをラップしますが、問題が発生した場合に備えて、その存在を認識しておくことをお勧めします。また、バッチが参加するランク間でシャード化される前に、バッチ辞書に正しい shift_labels を自動的に注入します。

ALST/UlyssesSPの使用を開始するために残っている唯一の要素は、勾配を正しく取得するために、異なる微分可能なall_gatherを使用してランク間で損失を集約することです。以下のコードは、-100でマスクされたトークンを除外しながらこれを行い、正しい平均値を取得します。

sp_size = parallelism_config.sp_size if parallelism_config is not None else 1
if sp_size > 1:
    sp_group = accelerator.torch_device_mesh["sp"].get_group()
    sp_world_size = parallelism_config.sp_size

# Normal training loop
for iter, batch in enumerate(dl):
    optimizer.zero_grad()

    batch = move_to_device(batch, model.device)

    # The model automatically receives shift_labels via **kwargs and uses it for loss computation.
    # Both standard transformers models and Liger-patched models handle this correctly.
    outputs = model(**batch)
    loss = outputs.loss
    shift_labels = batch["shift_labels"]

    if sp_size > 1:
        # differentiable weighted per-shard-loss aggregation across ranks
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (shift_labels != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(
            good_tokens, group=sp_group
        )
        # Skip ranks with zero valid tokens to avoid NaN contamination (NaN * 0 = NaN)
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

    if rank == 0: accelerator.print(f"{iter}: {loss=}")
    accelerator.log(dict(train_loss=loss, step=iter))

    accelerator.backward(loss)
    optimizer.step()

shift_labels がバッチに存在する場合、モデルはこれを自動的に処理する点に注意してください。モデルのフォワードパスは**kwargs を介してshift_labels を受け取り、損失関数に渡します。これにより、シーケンス並列性に対する損失が正しく計算されます。Liger Kernel を使用する場合も、shift_labels をシームレスに処理し、非常にメモリ効率の高い方法で損失を計算します。Liger は長いシーケンス長に対して強く推奨されます。これは、融合演算（メモリに完全な logits テンソルをマテリアライズしない融合 logits-loss 計算など）を使用することで GPU メモリを解放するからです。

HF Accelerate が裏でどのような処理を行ったかを確認したい場合は、この完全な統合チュートリアルをお読みください。

ALST/UlyssesSPを有効にしたAccelerateのトレーニングループの例については、examples/alst_ulysses_sequence_parallelismを参照してください。

[!Warning]

このAPIは非常に新しく、まだ実験段階にあります。安定したAPIの提供に努めていますが、パブリックAPIの一部が将来変更される可能性があります。

これはDeepspeedバックエンドであるため、通常のDeepspeed設定が適用されます。そのため、シーケンス並列性をオプティマイザ状態および/または重みのオフロードと組み合わせることで、より多くのGPUメモリを解放し、さらに長いシーケンス長を有効にできます。この技術は、DeepSpeed ZeROステージ2および3で動作することがテストされています。