diffusers/main/ja/optimization/context_parallelism.md

コンテキスト並列

Context parallelismは入力シーケンスをGPU間で分割し、各GPUが自身のスライスを処理します。これによりGPUごとのメモリ使用量が削減され、高解像度画像や動画のような長いシーケンスにおいて特に効果的です。他の戦略との比較については、Distributed inferenceを参照してください。

Context parallelism は attention バックエンド内で実行され、サポートしているバックエンドは一部のみです。お使いのバックエンドがサポートしていない場合、enable_parallelism は互換性のあるバックエンドを列挙したエラーを発生させます。set_attention_backend() を使用してバックエンドを切り替え、利用可能なバックエンド で完全なリストを確認してください。

以下のすべてのバリアントを ContextParallelConfig で設定します。

バリアント	設定	要件	使用場面
リングアテンション	ring_degree > 1	シーケンス長はring_degreeで割り切れる必要があります	シーケンスが長く、インターコネクト帯域幅が制限されている
Ulysses Attention	ulysses_degree > 1	シーケンス長とアテンションヘッド数はulysses_degreeで割り切れる必要があります	GPUが高速なインターコネクトを備えている
Unified Attention	ring_degree > 1かつulysses_degree > 1	4 以上の GPU	Ulysses Attention単独で使用できるデバイス数よりも、利用可能なデバイス数が多い
Ulysses Anything Attention	ulysses_degree > 1, ulysses_anything=True	ring_degree は 1 である必要があります	シーケンス長またはヘッド数が割り切れない場合
Ring Anything Attention	ring_degree > 1, ring_anything=True	ulysses_degree は 1 である必要があります。推論専用です。attn_mask は None である必要があります。	シーケンス長が均等に割り切れない

リングアテンション

Ring Attentionは、リング状に配置されたGPU間でキー（K）とバリュー（V）のブロックを伝達します。各GPUは自身が保持するK/Vブロックに対してアテンションを計算し、その後そのブロックを次のGPUに渡します。1周のループが完了すると、すべてのGPUが全トークンのK/Vを確認したことになりますが、どのGPUも完全なシーケンスを保持することはありません。

ring_degree を GPU の数に設定し、その設定を Transformer の enable_parallelism() メソッドに渡します。

import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    # torchrun sets LOCAL_RANK to this process's GPU index on its node
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)
    return device

def main():
    device = setup_distributed()
    world_size = dist.get_world_size()

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
    ).to(device)
    pipeline.transformer.set_attention_backend("_native_cudnn")

    cp_config = ContextParallelConfig(ring_degree=world_size)
    pipeline.transformer.enable_parallelism(config=cp_config)

    prompt = """
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
    """

    # Must specify generator so all ranks start with same latents (or pass your own)
    generator = torch.Generator().manual_seed(42)
    image = pipeline(
        prompt,
        guidance_scale=3.5,
        num_inference_steps=50,
        generator=generator,
    ).images[0]

    if dist.get_rank() == 0:
        image.save("output.png")

    if dist.is_initialized():
        dist.destroy_process_group()

if __name__ == "__main__":
    main()

スクリプトをring_attention.pyとして保存し、torchrunまたは他のPyTorch互換のランチャーで起動します。--nproc-per-nodeにはGPUの数を指定します。

torchrun --nproc-per-node 2 ring_attention.py

他のバリアントを試すには、ring_attention.py 内の enable_parallelism の呼び出しを、それぞれのセクションにあるものに置き換えてください。

Ulysses Attention

Ulysses Attention uses all-to-all communication, where every GPU exchanges data with every other GPU. The first all-to-all regroups the data so each GPU holds every token for a subset of attention heads. Each GPU computes attention for its heads, and a second all-to-all regroups the results by token for the next layer.

ulysses_degree を GPU の数に設定し、その設定を enable_parallelism() に渡します。

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

統一アテンション

Unified Sequence Parallelismは、Ring AttentionとUlysses Attentionを組み合わせたものです。all-to-allを実行してヘッドとトークンを再分配し、その結果に対してRing Attentionを実行してから、all-to-allを逆方向に実行して元のレイアウトを復元します。Ulyssesはアテンションヘッド間で並列化し、Ringはシーケンス間で並列化するため、両方を組み合わせることで、Ulysses単独よりも多くのGPUを利用できます。

GPU が 4 台以上ある場合は Unified Attention を使用してください。ulysses_degree と ring_degree の両方を 1 より大きい値に設定します。GPU は ulysses_degree * ring_degree 台のデバイスからなる 2 次元グリッドを形成し、各 GPU は 1 つの Ulysses グループと 1 つの Ring グループに属します。

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2))

以下のベンチマークでは、このスクリプトを使用して、4つのH100 GPU上でUlysses、Ring、Unified Attentionを比較しています。各画像は50回の推論ステップを使用しています。

CP バックエンド	時間 / 画像 (ms)	ステップ / 秒	ピークメモリ (GB)
ulysses	6670.789	7.50	33.85
リング	13076.492	3.82	56.02
unified_balanced	11068.705	4.52	33.85

Ulyssesはスループットが最も高いですが、モデルのattention headsの数を超えるGPUを使用できません。Unified Attentionはこの制限を解消します。

Ulysses Anything Attention

Ulysses Attention では、シーケンス長とアテンションヘッド数がulysses_degreeで割り切れる必要があります。Ulysses Anything Attentionは、これらの両方の要件を解消します。Unified Attention と組み合わせることはできないため、ring_degreeは 1 でなければなりません。

ulysses_degreeを1より大きい値に設定し、ulysses_anything=Trueにします。

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))

init_process_group に CPU テンソル用の gloo バックエンドを追加します。これにより、ホストからデバイスへの転送やデバイスからホストへの転送によって引き起こされる強制的な CUDA 同期を回避できます。

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

以下のベンチマークは、このスクリプトを使用して、4つのL20 GPU上でFLUX.1-devを実行します。

CP バックエンド	時間 / イテレーション (ms)	ステップ / 秒	ピークメモリ (GB)	形状 (HxW)
ulysses	281.07	3.56	37.11	1024x1024
リング	351.34	2.85	37.01	1024x1024
unified_balanced	324.37	3.08	37.16	1024x1024
ulysses_anything	280.94	3.56	37.11	1024x1024
ulysses	失敗	失敗	失敗	1008x1008
リング	失敗	失敗	失敗	1008x1008
unified_balanced	失敗	失敗	失敗	1008x1008
ulysses_anything	278.40	3.59	36.99	1008x1008

1008x1008では、シーケンス長が均等に割り切れないため、Ulysses Anything Attentionのみが実行されます。1024x1024では、標準のUlysses Attentionと同等のパフォーマンスを発揮します。

Ring Anything Attention

Ring Attention は、シーケンス長が ring_degree で割り切れることを必要とします。Ring Anything Attention は、任意のシーケンス長をサポートします。各 GPU のローカル K/V を最長のローカル長までパディングし、パディングされた K/V を all-gather してから、アテンションを計算する前に各ブロックを本来の長さにスライスします。

Unified Attention と組み合わせることはできないため、ulysses_degree は 1 でなければなりません。ring_degree を 1 以上に設定し、ring_anything=True を指定してください。

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))

Ring Anything Attention は、Ulysses Anything Attention に示されている gloo バックエンドの恩恵も受けます。

Ring Anything Attention は推論のみをサポートしており、attn_mask は None である必要があります。

以下のベンチマークは、4つのRTX 4090 (48GB) GPU上でFLUX.1-devを実行します。

CP バックエンド	時間 / イテレーション (ms)	ステップ / 秒	ピークメモリ (GB)	形状 (HxW)
ulysses	259.07	3.86	33.83	1024x1024
リング	338.98	2.95	33.83	1024x1024
unified_balanced	321.54	3.11	33.83	1024x1024
ulysses_anything	259.07	3.86	33.83	1024x1024
ring_anything	340.14	2.94	33.83	1024x1024
ulysses	失敗	失敗	失敗	1008x1008
リング	失敗	失敗	失敗	1008x1008
unified_balanced	失敗	失敗	失敗	1008x1008
ulysses_anything	253.16	3.95	33.75	1008x1008
ring_anything	335.57	2.98	33.75	1008x1008

Ring Anything Attention は 1008x1008 で実行され、標準の Ring Attention とほぼ同じパフォーマンスを発揮します。動機とトレードオフについては、このコメントを参照してください。

…

enable_parallelismをロード後に呼び出す代わりに、from_pretrained()のparallel_config引数にContextParallelConfigを渡します。まずプロセスグループを初期化してください。from_pretrainedは、torch.distributedが初期化されていない場合にエラーを発生させるためです。Ring Attentionの例と同様に、スクリプトはtorchrunで起動します。

import os
import torch
from torch import distributed as dist
from diffusers import AutoModel, ContextParallelConfig, DiffusionPipeline

CKPT_ID = "black-forest-labs/FLUX.1-dev"

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)

cp_config = ContextParallelConfig(ring_degree=dist.get_world_size())
transformer = AutoModel.from_pretrained(
    CKPT_ID,
    subfolder="transformer",
    dtype=torch.bfloat16,
    parallel_config=cp_config
)

pipeline = DiffusionPipeline.from_pretrained(
    CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
).to(device)