コンテキスト並列

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

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

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

リングアテンション

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回の推論ステップを使用しています。

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を実行します。

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を実行します。

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)