diffusers/pr_14039/en/api/parallel.md

Parallelism

Parallelism strategies help speed up diffusion transformers by distributing computations across multiple devices, allowing for faster inference/training times. Refer to the Distributed inferece guide to learn more.

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)

Source

Parameters:

context_parallel_config (ContextParallelConfig, optional) : Configuration for context parallelism.

tensor_parallel_config (TensorParallelConfig, optional) : Configuration for tensor parallelism.

Configuration for applying different parallelisms.

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)

Source

Parameters:

ring_degree (int, optional, defaults to 1) : Number of devices to use for Ring Attention. Sequence is split across devices. Each device computes attention between its local Q and KV chunks passed sequentially around ring. Lower memory (only holds 1/N of KV at a time), overlaps compute with communication, but requires N iterations to see all tokens. Best for long sequences with limited memory/bandwidth. Number of devices to use for ring attention within a context parallel region. Must be a divisor of the total number of devices in the context parallel mesh.

ulysses_degree (int, optional, defaults to 1) : Number of devices to use for Ulysses Attention. Sequence split is across devices. Each device computes local QKV, then all-gathers all KV chunks to compute full attention in one pass. Higher memory (stores all KV), requires high-bandwidth all-to-all communication, but lower latency. Best for moderate sequences with good interconnect bandwidth.

convert_to_fp32 (bool, optional, defaults to True) : Whether to convert output and LSE to float32 for ring attention numerical stability.

rotate_method (str, optional, defaults to "allgather") : Method to use for rotating key/value states across devices in ring attention. Currently, only "allgather" is supported.

ulysses_anything (bool, optional, defaults to False) : Whether to enable "Ulysses Anything" mode, which supports arbitrary sequence lengths and head counts that are not evenly divisible by ulysses_degree. When enabled, ulysses_degree must be greater than 1 and ring_degree must be 1.

ring_anything (bool, optional, defaults to False) : Whether to enable "Ring Anything" mode, which supports arbitrary sequence lengths. When enabled, ring_degree must be greater than 1 and ulysses_degree must be 1.

mesh (torch.distributed.device_mesh.DeviceMesh, optional) : A custom device mesh to use for context parallelism. If provided, this mesh will be used instead of creating a new one. This is useful when combining context parallelism with other parallelism strategies (e.g., FSDP, tensor parallelism) that share the same device mesh. The mesh must have both "ring" and "ulysses" dimensions. Use size 1 for dimensions not being used (e.g., mesh_shape=(2, 1, 4) with mesh_dim_names=("ring", "ulysses", "fsdp") for ring attention only with FSDP).

Configuration for context parallelism.

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

diffusers.hooks.apply_context_parallel(module: Module, parallel_config: ContextParallelConfig, plan: dict)

Source

Apply context parallel on a model.

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)

Source

Parameters:

tp_degree (int, defaults to 1) : Number of devices to shard across. Must be a divisor of the number of attention heads (and FFN hidden dimensions) of the model being parallelised.

mesh (torch.distributed.device_mesh.DeviceMesh, optional) : A custom device mesh to use. If provided, tp_degree is inferred from mesh.size() and the argument is ignored. Useful when combining TP with other parallelism strategies (e.g. CP) that share the same mesh.

Configuration for tensor parallelism.

Tensor parallelism shards weight matrices (column-wise and row-wise) across devices. Each device computes a partial result; an AllReduce/AllGather at layer boundaries reconstructs the full output. Uses torch.distributed.tensor.parallelize_module with ColwiseParallel / RowwiseParallel sharding styles. Supported device types are "cuda" and "neuron".

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

diffusers.hooks.apply_tensor_parallel(model: Module, config: TensorParallelConfig, tp_plan: dict)

Source

Apply tensor parallel on a model from its flat _tp_plan.

TPU

enable_tpu_compile[[diffusers.DiffusionPipeline.enable_tpu_compile]]

diffusers.DiffusionPipeline.enable_tpu_compile[[diffusers.DiffusionPipeline.enable_tpu_compile]]

diffusers.DiffusionPipeline.enable_tpu_compile(model_names: typing.Optional[typing.List[str]] = None, **compile_kwargs)

Source

Parameters:

model_names (list[str], optional) : Names of pipeline components to compile. Defaults to all torch.nn.Module components currently resident on a TPU device.

• **compile_kwargs : Extra keyword arguments forwarded to torch.compile. backend defaults to TpuBackend() and dynamic defaults to False (required for static tracing).

Compile pipeline components that are on TPU using torch.compile with the TpuBackend.

Before compiling, each component that exposes set_attn_processor has AttnProcessor applied. This replaces AttnProcessor2_0 (SDP-based) which triggers XLA fusion-emitter crashes in eager/lazy mode. TpuBackend handles the resulting torch.cat layout internally during static tracing, so no additional wrapper is needed at compile time.

Example:

import torch
import torch_tpu  # noqa: F401

pipe.transformer.to("tpu")
pipe.vae.to("tpu")
pipe.enable_tpu_compile()

tpu_warmup[[diffusers.DiffusionPipeline.tpu_warmup]]

diffusers.DiffusionPipeline.tpu_warmup[[diffusers.DiffusionPipeline.tpu_warmup]]

diffusers.DiffusionPipeline.tpu_warmup(*args, **kwargs)

Source

Parameters:

• *args : Positional arguments forwarded to the pipeline __call__.

• **kwargs : Keyword arguments forwarded to the pipeline __call__.

Run a single forward pass to trigger XLA / TpuBackend compilation.

Call this after enable_tpu_compile and before timed inference. The warmup pass compiles the static computation graphs; subsequent calls reuse the compiled graphs and run at full speed.

Example:

pipe.tpu_warmup(
    prompt="warmup",
    height=1024,
    width=1024,
    num_inference_steps=4,
    guidance_scale=0.0,
)