# 上下文并行

[上下文并行](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism)将输入序列拆分到多个 GPU 上，每个 GPU 处理其对应的分片。这降低了每个 GPU 的内存占用，对于高分辨率图像和视频等长序列尤为有益。要将其与其他策略进行比较，请参阅[分布式推理](../training/distributed_inference)。

上下文并行在注意力后端内部运行，且只有部分后端支持该功能。如果您的后端不支持，`enable_parallelism` 将抛出一个错误，其中列出了兼容的后端。使用 [set_attention_backend()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.set_attention_backend) 切换后端，并参阅 [可用后端](./attention_backends#available-backends) 以获取完整列表。

使用 [ContextParallelConfig](/docs/diffusers/v0.43.0/zh/api/parallel#diffusers.ContextParallelConfig) 配置以下每个变体。

| 变体 | 配置 | 要求 | 使用时 |
|---|---|---|---|
| [环形注意力](#ring-attention) | `ring_degree > 1` | 序列长度必须能被`ring_degree`整除 | 序列较长且互连带宽有限 |
| [Ulysses 注意力](#ulysses-attention) | `ulysses_degree > 1` | 序列长度和注意力头数必须能被`ulysses_degree`整除 | 您的 GPU 配备了高速互连 |
| [统一注意力](#unified-attention) | `ring_degree > 1` 且 `ulysses_degree > 1` | 至少 4 个 GPU | 你拥有的设备数量超过了 Ulysses Attention 单独能够使用的数量 |
| [Ulysses 任意注意力](#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。经过一轮完整的循环后，每个 GPU 都已处理过所有 token 的 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 注意力

[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 以重新分配 heads 和 tokens，然后对结果运行 Ring Attention，最后反向执行 all-to-all 以恢复原始布局。Ulysses 在 attention heads 维度上进行并行，而 Ring 在序列维度上进行并行，因此两者结合后可以比单独使用 Ulysses 利用更多的 GPU。

当拥有至少 4 个 GPU 时，请使用 Unified Attention。将 `ulysses_degree` 和 `ring_degree` 均设置为大于 1 的值。这些 GPU 会组成一个由 `ulysses_degree * ring_degree` 个设备构成的二维网格，每个 GPU 属于一个 Ulysses 组和一个 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 后端         | 时间 / 图像 (毫秒) | 步数 / 秒 | 峰值内存（GB） |
|--------------------|------------------|-------------|------------------|
| ulysses            | 6670.789         | 7.50        | 33.85            |
| ring               | 13076.492        | 3.82        | 56.02            |
| unified_balanced   | 11068.705        | 4.52        | 33.85            |

Ulysses 的吞吐量最高，但它使用的 GPU 数量不能超过模型的注意力头数。Unified Attention 移除了这一限制。

## Ulysses 任意注意力

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

以下基准测试在 4 块 L20 GPU 上使用[此脚本](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999)运行 FLUX.1-dev。

| CP 后端         | 时间 / 迭代次数（毫秒） | 步数 / 秒 | 峰值内存（GB） | 形状 (HxW)|
|--------------------|------------------|-------------|------------------|------------|
| ulysses            |   281.07         |    3.56     |     37.11        | 1024x1024  |
| ring               |   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  |
| ring               |   失败         |    失败   |     失败       | 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 后端         | 时间 / 迭代次数（毫秒） | 步数 / 秒 | 峰值内存（GB） | 形状 (HxW)|
|--------------------|------------------|-------------|------------------|------------|
| ulysses            |   259.07         |    3.86     |     33.83        | 1024x1024  |
| ring               |   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  |
| ring               |   失败         |    失败   |     失败       | 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`，不如将 [ContextParallelConfig](/docs/diffusers/v0.43.0/zh/api/parallel#diffusers.ContextParallelConfig) 传递给 [from_pretrained()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.from_pretrained) 的 `parallel_config` 参数。请先初始化进程组，因为如果 `torch.distributed` 未初始化，`from_pretrained` 会引发错误。请像 [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)
```