diffusers/v0.43.0/zh/api/parallel.md

并行

并行策略通过将计算分布在多个设备上，帮助加速扩散 Transformer，从而实现更快的推理/训练速度。有关更多信息，请参阅分布式推理指南。

ParallelConfig[[diffusers.ParallelConfig]]

diffusers.ParallelConfig[[diffusers.ParallelConfig]]

diffusers.ParallelConfig(context_parallel_config: diffusers.models._modeling_parallel.ContextParallelConfig | None = None, tensor_parallel_config: diffusers.models._modeling_parallel.TensorParallelConfig | None = None, _rank: int = None, _world_size: int = None, _device: device = None, _mesh: DeviceMesh = None)

源代码

参数：

context_parallel_config （ContextParallelConfig，可选）：上下文并行的配置。

tensor_parallel_config (TensorParallelConfig，可选)：张量并行的配置。

用于应用不同并行策略的配置。

ContextParallelConfig[[diffusers.ContextParallelConfig]]

diffusers.ContextParallelConfig[[diffusers.ContextParallelConfig]]

diffusers.ContextParallelConfig(ring_degree: int | None = None, ulysses_degree: int | None = None, convert_to_fp32: bool = True, rotate_method: typing.Literal['allgather', 'alltoall'] = 'allgather', mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None, ulysses_anything: bool = False, ring_anything: bool = False, _rank: int = None, _world_size: int = None, _device: device = None, _mesh: DeviceMesh = None, _flattened_mesh: DeviceMesh = None, _ring_mesh: DeviceMesh = None, _ulysses_mesh: DeviceMesh = None, _ring_local_rank: int = None, _ulysses_local_rank: int = None)

源代码

参数：

ring_degree (int, 可选，默认为 1)：用于 Ring Attention 的设备数量。序列将在设备之间进行拆分。每个设备计算其本地 Q 与按顺序在环中传递的 KV 块之间的注意力。内存占用更低（每次仅持有 1/N 的 KV），计算与通信可重叠，但需要 N 次迭代才能看到所有 token。最适合内存/带宽受限的长序列场景。在上下文并行区域内用于 ring attention 的设备数量。必须是上下文并行网格中设备总数的约数。

ulysses_degree (int，可选，默认为 1)：用于 Ulysses Attention 的设备数量。序列将在设备之间进行拆分。每个设备计算本地的 QKV，然后通过 all-gather 收集所有 KV 块，以一次性计算完整的注意力。这种方式内存占用较高（需存储所有 KV），且需要高带宽的 all-to-all 通信，但延迟较低。适用于具有良好互连带宽的中等长度序列。

convert_to_fp32 (bool，可选，默认为 True)：是否将输出和 LSE 转换为 float32，以确保 ring attention 的数值稳定性。

rotate_method (str，可选，默认为 "allgather")：在环形注意力中用于跨设备轮换键/值状态的方法。目前仅支持 "allgather"。

ulysses_anything (bool，可选，默认为 False)：是否启用“Ulysses Anything”模式。该模式支持无法被 ulysses_degree 整除的任意序列长度和头数。启用此模式时，ulysses_degree 必须大于 1，且 ring_degree 必须为 1。

ring_anything (bool，可选，默认为 False)：是否启用“Ring Anything”模式，该模式支持任意序列长度。启用后，ring_degree 必须大于 1，且 ulysses_degree 必须为 1。

mesh (torch.distributed.device_mesh.DeviceMesh, 可选) : 用于上下文并行的自定义设备网格。如果提供，将使用该网格而不是创建新的网格。当将上下文并行与其他共享同一设备网格的并行策略（例如 FSDP、张量并行）结合使用时，这非常有用。该网格必须同时具有 "ring" 和 "ulysses" 维度。对于未使用的维度，请使用大小 1（例如，仅使用 FSDP 的 ring attention 时，使用 mesh_shape=(2, 1, 4) 和 mesh_dim_names=("ring", "ulysses", "fsdp")）。

用于上下文并行的配置。

diffusers.hooks.apply_context_parallel[[diffusers.hooks.apply_context_parallel]]

diffusers.hooks.apply_context_parallel(module: Module, parallel_config: ContextParallelConfig, plan: dict[str, dict[str, dict[str | int, diffusers.models._modeling_parallel.ContextParallelInput | list[diffusers.models._modeling_parallel.ContextParallelInput] | tuple[diffusers.models._modeling_parallel.ContextParallelInput, ...]] | diffusers.models._modeling_parallel.ContextParallelOutput | list[diffusers.models._modeling_parallel.ContextParallelOutput] | tuple[diffusers.models._modeling_parallel.ContextParallelOutput, ...]]])

源代码

在模型上应用上下文并行。

TensorParallelConfig[[diffusers.TensorParallelConfig]]

diffusers.TensorParallelConfig[[diffusers.TensorParallelConfig]]

diffusers.TensorParallelConfig(tp_degree: int = 1, mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None, _rank: int = None, _world_size: int = None, _device: device = None, _mesh: DeviceMesh = None, _tp_degree: int = None)

源代码

参数：

tp_degree (int，默认为 1) : 用于分片的设备数量。必须是正在并行化的模型的注意力头数（以及 FFN 隐藏维度）的约数。

mesh (torch.distributed.device_mesh.DeviceMesh, 可选) : 要使用的自定义设备网格。如果提供，tp_degree 将从 mesh.size() 推断，该参数将被忽略。在将 TP 与其他共享同一网格的并行策略（例如 CP）结合使用时非常有用。

用于张量并行的配置。

张量并行将权重矩阵（按列和按行）分片到多个设备上。每个设备计算部分结果；在层边界处通过 AllReduce/AllGather 重建完整输出。使用 torch.distributed.tensor.parallelize_module 配合 ColwiseParallel / RowwiseParallel 分片方式。支持的设备类型为 "cuda"、"neuron" 和 "tpu"。

diffusers.hooks.apply_tensor_parallel[[diffusers.hooks.apply_tensor_parallel]]

diffusers.hooks.apply_tensor_parallel(model: Module, config: TensorParallelConfig, tp_plan: dict, weights_already_sharded: bool = False)

源代码

参数：

model (torch.nn.Module)：要就地分片的模型。

config (TensorParallelConfig)：张量并行配置。其设备网格必须已通过 config.setup(...) 完成设置。

tp_plan (dict) : 一个扁平映射，将模块名称通配符映射到 "colwise"/"rowwise" 风格（或其打包变体），通常是模型的 _tp_plan。

weights_already_sharded (bool，默认为 False)：计划中的参数是否已经是 DTensor 分片，例如在流式 from_pretrained 加载之后。如果为 True，则仅注册前向钩子。之所以显式传入该值而非自动检测，是因为如果某个计划中的参数缺失于检查点中，它仍然会是一个 meta 张量，这会导致检测逻辑将一个实际上已部分分片的模型误判为“未分片”。

根据模型的扁平化 _tp_plan 对其应用张量并行。