# コンテキスト並列

[Context parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism)は入力シーケンスをGPU間で分割し、各GPUが自身のスライスを処理します。これによりGPUごとのメモリ使用量が削減され、高解像度画像や動画のような長いシーケンスにおいて特に効果的です。他の戦略との比較については、[Distributed inference](../training/distributed_inference)を参照してください。

Context parallelism は attention バックエンド内で実行され、サポートしているバックエンドは一部のみです。お使いのバックエンドがサポートしていない場合、`enable_parallelism` は互換性のあるバックエンドを列挙したエラーを発生させます。[set_attention_backend()](/docs/diffusers/v0.43.0/ja/api/models/overview#diffusers.ModelMixin.set_attention_backend) を使用してバックエンドを切り替え、[利用可能なバックエンド](./attention_backends#available-backends) で完全なリストを確認してください。

以下のすべてのバリアントを [ContextParallelConfig](/docs/diffusers/v0.43.0/ja/api/parallel#diffusers.ContextParallelConfig) で設定します。

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

## リングアテンション

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

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

```py
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](https://docs.pytorch.org/docs/stable/elastic/run.html)または他のPyTorch互換のランチャーで起動します。`--nproc-per-node`にはGPUの数を指定します。

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

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

## Ulysses Attention

[Ulysses Attention](https://huggingface.co/papers/2309.14509) 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()` に渡します。

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

## 統一アテンション

[Unified Sequence Parallelism](https://huggingface.co/papers/2405.07719)は、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 グループに属します。

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

以下のベンチマークでは、[このスクリプト](https://github.com/huggingface/diffusers/pull/12693#issuecomment-3694727532)を使用して、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](https://github.com/huggingface/diffusers/pull/12996)は、これらの両方の要件を解消します。Unified Attention と組み合わせることはできないため、`ring_degree`は 1 でなければなりません。

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

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

> [!TIP]
> [init_process_group](https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) に CPU テンソル用の `gloo` バックエンドを追加します。これにより、ホストからデバイスへの転送やデバイスからホストへの転送によって引き起こされる強制的な CUDA 同期を回避できます。

```py
import torch.distributed as dist

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

以下のベンチマークは、[このスクリプト](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999)を使用して、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](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4302195582) は、任意のシーケンス長をサポートします。各 GPU のローカル K/V を最長のローカル長までパディングし、パディングされた K/V を all-gather してから、アテンションを計算する前に各ブロックを本来の長さにスライスします。

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

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

Ring Anything Attention は、[Ulysses Anything Attention](#ulysses-anything-attention) に示されている `gloo` バックエンドの恩恵も受けます。

> [!NOTE]
> 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 とほぼ同じパフォーマンスを発揮します。動機とトレードオフについては、[このコメント](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4304104462)を参照してください。

## …

`enable_parallelism`をロード後に呼び出す代わりに、[from_pretrained()](/docs/diffusers/v0.43.0/ja/api/models/overview#diffusers.ModelMixin.from_pretrained)の`parallel_config`引数に[ContextParallelConfig](/docs/diffusers/v0.43.0/ja/api/parallel#diffusers.ContextParallelConfig)を渡します。まずプロセスグループを初期化してください。`from_pretrained`は、`torch.distributed`が初期化されていない場合にエラーを発生させるためです。[Ring Attention](#ring-attention)の例と同様に、スクリプトは`torchrun`で起動します。

```py
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)
```