上下文并行

上下文并行将输入序列拆分到多个 GPU 上，每个 GPU 处理其对应的分片。这降低了每个 GPU 的内存占用，对于高分辨率图像和视频等长序列尤为有益。要将其与其他策略进行比较，请参阅分布式推理。

上下文并行在注意力后端内部运行，且只有部分后端支持该功能。如果您的后端不支持，enable_parallelism 将抛出一个错误，其中列出了兼容的后端。使用 set_attention_backend() 切换后端，并参阅可用后端以获取完整列表。

使用 ContextParallelConfig 配置以下每个变体。

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

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 以重新分配 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 组。

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

下方的基准测试使用此脚本在 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 移除了这两项要求。它不能与 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 后端	时间 / 迭代次数（毫秒）	步数 / 秒	峰值内存（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 支持任意序列长度。它会将每个 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 后端	时间 / 迭代次数（毫秒）	步数 / 秒	峰值内存（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 大致相当。有关动机和权衡，请参阅此评论。

在加载时启用上下文并行

与其在加载后调用 enable_parallelism，不如将 ContextParallelConfig 传递给 from_pretrained() 的 parallel_config 参数。请先初始化进程组，因为如果 torch.distributed 未初始化，from_pretrained 会引发错误。请像 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)

Update on GitHub

Diffusers

上下文并行

环形注意力

Ulysses 注意力

统一注意力

Ulysses 任意注意力

Ring Anything Attention

在加载时启用上下文并行