Buckets:
| # 上下文并行 | |
| [上下文并行](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) | |
| ``` | |
Xet Storage Details
- Size:
- 12.1 kB
- Xet hash:
- da7f9db144b411b52f873175bf7caf380656a832d5ebb404564b77c7371fe99f
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.