Allow tensor parallelism and context parallelism in one ParallelConfig

docs/source/en/training/distributed_inference.md

@@ -561,6 +561,79 @@ Validate a new plan numerically rather than by eye: generate with a fixed seed o
 > [!TIP]
 > Start from an existing plan for a similar architecture. [`QwenImageTransformer2DModel`] is fully unfused and every entry is plain `"colwise"`/`"rowwise"`, [`FluxTransformer2DModel`] adds a single packed row-wise projection, and [`Flux2Transformer2DModel`] covers both packed styles — see [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py), [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py), and [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) respectively.
 
+## Combining context and tensor parallelism
+
+The two strategies split a model along different axes: tensor parallelism shards the weights, context parallelism splits the sequence. A model whose weights don't fit on one device *and* whose sequence is too long for one device's attention needs both, and neither alone can help: TP doesn't divide the sequence, and CP doesn't divide the weights.
+
+Set both in one [`ParallelConfig`]. [`~ModelMixin.enable_parallelism`] (or `from_pretrained(..., parallel_config=...)`) builds a single device mesh with a dimension for each, `("tp", "ring", "ulysses")`, so the TP all-reduces and the CP all-gathers or all-to-alls each stay inside their own process group. The model must define both a `_tp_plan` and a `_cp_plan`.
+
+```py
+import torch
+from torch import distributed as dist
+from diffusers import (
+    ContextParallelConfig,
+    DiffusionPipeline,
+    Flux2Transformer2DModel,
+    ParallelConfig,
+    TensorParallelConfig,
+)
+
+def main():
+    dist.init_process_group(backend="nccl")
+    rank = dist.get_rank()
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    # 4 devices: each weight is sharded 2 ways, and each shard's sequence is split 2 ways with Ulysses.
+    transformer = Flux2Transformer2DModel.from_pretrained(
+        "black-forest-labs/FLUX.2-dev",
+        subfolder="transformer",
+        torch_dtype=torch.bfloat16,
+        parallel_config=ParallelConfig(
+            tensor_parallel_config=TensorParallelConfig(tp_degree=2),
+            context_parallel_config=ContextParallelConfig(ulysses_degree=2),
+        ),
+    )
+
+    pipeline = DiffusionPipeline.from_pretrained(
+        "black-forest-labs/FLUX.2-dev", transformer=transformer, torch_dtype=torch.bfloat16
+    )
+    # As with tensor parallelism alone, move the other components individually, not with `pipeline.to(device)`.
+    pipeline.text_encoder.to(device)
+    pipeline.vae.to(device)
+
+    generator = torch.Generator().manual_seed(42)
+    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
+    if rank == 0:
+        image.save("output.png")
+    dist.destroy_process_group()
+
+if __name__ == "__main__":
+    main()
+```
+
+```shell
+torchrun --nproc-per-node 4 tp_cp_flux2.py
+```
+
+The degrees are checked before anything is sharded:
+
+- `tp_degree` × `ring_degree` × `ulysses_degree` must not exceed the world size.
+- Ulysses splits the heads that TP left on each rank, so `num_attention_heads` must be divisible by `tp_degree` × `ulysses_degree`. Ring attention doesn't touch the head dimension, so with `ring_degree` only `tp_degree` has to divide the head count. Pass `ulysses_anything=True` to pad the heads instead when they don't divide.
+
+The default mesh puts `"tp"` outermost, so the Ulysses groups are runs of consecutive ranks. Some accelerator runtimes, AWS Neuron among them, only accept contiguous replica groups for all-to-all while tolerating strided groups for the all-reduce TP needs, and only one axis of the mesh can be contiguous. On multi-node CUDA the opposite order is usually better, since TP's per-layer all-reduce is the most bandwidth-hungry collective and wants to stay inside one NVLink domain. To choose the layout yourself, build the mesh with `"tp"`, `"ring"` and `"ulysses"` dimensions and pass it as `mesh=` on either config.
+
+Which degree to raise first depends on what you are short of. TP is what lowers per-device memory; CP mostly lowers latency. For example, [MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) (a 33B audio-video DiT, 56 heads) at 1344×768 and 124 frames on one AWS Trainium2 instance (trn2.48xlarge, loaded with `from_pretrained(..., parallel_config=...)`; peak memory is measured during denoising):
+
+| NeuronCores | config | weights / core | peak memory / core | s / step |
+|---|---|---|---|---|
+| 8 | TP=8 | 7.84 GiB | 18.7 GiB | 9.33 |
+| 16 | TP=4 × Ulysses=4 | 15.53 GiB | 21.7 GiB | 5.31 |
+| 32 | TP=4 × Ulysses=8 | 15.53 GiB | 20.5 GiB | 3.97 |
+| 32 | TP=8 × Ulysses=4 | 7.84 GiB | 13.7 GiB | 3.81 |
+
+At the same 32 cores, TP=8 × Ulysses=4 is both faster and lighter than TP=4 × Ulysses=8, so on this model it pays to max out TP first and add CP on top. Its 56 heads cap TP at 8, so going further means raising the CP degree.
+
 ## Choosing a strategy
 
 The strategies above solve different problems, and the useful question is not which is fastest in the abstract but what you are running out of.
@@ -577,4 +650,4 @@ Some practical guidance:
 - **Throughput on many prompts, model already fits.** Use data parallelism. It is the only strategy here that scales throughput linearly without touching the model, and it leaves single-prompt latency alone.
 - **A single component's weights don't fit.** Reach for tensor parallelism first, since it lowers both memory and latency. It communicates at every block boundary, so it wants a fast interconnect like NVLink; over PCIe that per-layer traffic can outweigh the compute it saves, and `device_map` becomes the better choice. `device_map` also handles the case where the components are individually fine but collectively too large.
 - **Activations, not weights, are the problem.** This is the long-sequence regime — large images, many frames — and context parallelism is the direct answer. Pick a backend from the [Ulysses/Ring benchmarks](#ulysses-attention) above.
-- **Both weights and sequence are too large.** Combine tensor and context parallelism. [`TensorParallelConfig`] accepts a `mesh` argument so both can share one device mesh. This combination is not yet validated, so treat it as experimental.
+- **Both weights and sequence are too large.** [Combine tensor and context parallelism](#combining-context-and-tensor-parallelism) in one [`ParallelConfig`], raising `tp_degree` until the weights fit and adding context parallelism on top for latency.

src/diffusers/hooks/context_parallel.py

@@ -345,6 +345,17 @@ def _all_gather_anything(tensor: torch.Tensor, dim: int, group: dist.device_mesh
     rank_dim = shape[dim]
     gather_dims = gather_size_by_comm(rank_dim, group)
 
+    max_dim = max(gather_dims)
+    if any(d != max_dim for d in gather_dims) and tensor.device.type in ("neuron", "tpu"):
+        # Neuron and TPU all-gathers require every rank to send the same size, so pad each chunk to the longest one,
+        # gather, and trim each chunk back to its real length.
+        pad_shape = list(shape)
+        pad_shape[dim] = max_dim - rank_dim
+        padded = torch.cat([tensor, tensor.new_zeros(pad_shape)], dim=dim)
+        gathered_tensors = [torch.empty_like(padded) for _ in range(world_size)]
+        dist.all_gather(gathered_tensors, padded, group=group)
+        return torch.cat([t.narrow(dim, 0, d) for t, d in zip(gathered_tensors, gather_dims)], dim=dim)
+
     gather_shapes = _fill_gather_shapes(tuple(shape), tuple(gather_dims), dim, world_size)
 
     gathered_tensors = [torch.empty(shape, device=tensor.device, dtype=tensor.dtype) for shape in gather_shapes]

src/diffusers/models/_modeling_parallel.py

@@ -200,6 +200,12 @@ class ParallelConfig:
     """
     Configuration for applying different parallelisms.
 
+    Both may be set at once. The two are then applied over one device mesh with a dimension each — `("tp", "ring",
+    "ulysses")`, built by `enable_parallelism` — so their collectives stay in separate process groups. This is what
+    lets a model too large for one device (TP shards the weights) also run a sequence too long for one device's
+    attention scratchpad (CP splits the sequence): TP alone cannot divide the sequence, and CP alone cannot divide the
+    weights.
+
     Args:
         context_parallel_config (`ContextParallelConfig`, *optional*):
             Configuration for context parallelism.
@@ -216,12 +222,17 @@ class ParallelConfig:
     _mesh: torch.distributed.device_mesh.DeviceMesh = None
 
     def __post_init__(self):
-        if self.context_parallel_config is not None and self.tensor_parallel_config is not None:
+        if self.context_parallel_config is None and self.tensor_parallel_config is None:
             raise ValueError(
-                "Combining context parallelism and tensor parallelism in a single `ParallelConfig` is not supported. "
-                "Please specify only one of `context_parallel_config` or `tensor_parallel_config`."
+                "A `ParallelConfig` must specify at least one of `context_parallel_config` or "
+                "`tensor_parallel_config`."
             )
 
+    @property
+    def _is_combined(self) -> bool:
+        """Whether both context and tensor parallelism are requested, i.e. they must share one mesh."""
+        return self.context_parallel_config is not None and self.tensor_parallel_config is not None
+
     def setup(
         self,
         rank: int,
@@ -234,10 +245,32 @@ def setup(
         self._world_size = world_size
         self._device = device
         self._mesh = mesh
+
+        # Context and tensor parallelism compose because they cut the model along different axes: CP splits the
+        # sequence (and, under Ulysses, trades sequence for heads inside attention) while TP shards the Linear
+        # weights. Composing them means giving each its own mesh *dimension*, so that every collective one issues
+        # stays inside its own process group: TP's all-reduce must not reach a rank holding a different sequence
+        # chunk, and CP's sequence all-gather must not reach a rank holding a different weight shard.
+        cp_mesh = tp_mesh = mesh
+        if self._is_combined:
+            dim_names = mesh.mesh_dim_names if mesh is not None else None
+            missing = [d for d in ("ring", "ulysses", "tp") if dim_names is None or d not in dim_names]
+            if missing:
+                raise ValueError(
+                    f"Combining context and tensor parallelism requires a device mesh with 'ring', 'ulysses' and "
+                    f"'tp' dimensions, but {missing} {'is' if len(missing) == 1 else 'are'} missing (got "
+                    f"{dim_names}). `enable_parallelism` builds this mesh from `ring_degree`/`ulysses_degree`/"
+                    f"`tp_degree`; if you build it yourself, set `mesh=` on one of the two configs."
+                )
+            # `ContextParallelConfig.setup` slices ("ring", "ulysses") off whatever mesh it is handed, so it takes
+            # the full mesh. TP instead needs its own 1-D submesh: `parallelize_module` shards a weight over every
+            # rank of the mesh it is given, so handing it the 3-D mesh would shard over the CP ranks as well.
+            tp_mesh = mesh["tp"]
+
         if self.context_parallel_config is not None:
-            self.context_parallel_config.setup(rank, world_size, device, mesh)
+            self.context_parallel_config.setup(rank, world_size, device, cp_mesh)
         if self.tensor_parallel_config is not None:
-            self.tensor_parallel_config.setup(rank, world_size, device, mesh)
+            self.tensor_parallel_config.setup(rank, world_size, device, tp_mesh)
 
 
 @dataclass(frozen=True)

src/diffusers/models/modeling_utils.py

@@ -1546,6 +1546,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
                 # was recorded by `_resolve_parallel_config` before loading.
                 from ..hooks.tensor_parallel import apply_tensor_parallel
 
+                # A config that also asks for context parallelism gets its hooks here, in the same order as
+                # `enable_parallelism`, which is no longer callable now that the weights are sharded.
+                model._apply_context_parallel(parallel_config)
                 apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)
             else:
                 model.enable_parallelism(config=parallel_config)
@@ -1719,15 +1722,33 @@ def _resolve_parallel_config(
         device = torch.device(device_type, rank % torch.accelerator.device_count())
 
         mesh = None
-        if config.context_parallel_config is not None:
-            cp_config = config.context_parallel_config
+        cp_config, tp_config = config.context_parallel_config, config.tensor_parallel_config
+        if cp_config is not None and tp_config is not None:
+            # A single mesh spanning both parallelisms, with the context parallel dimensions varying fastest: a
+            # CP group is then a contiguous run of ranks and a TP group takes one rank out of each run.
+            #
+            # That order is a default, not a universal optimum. Some accelerator runtimes only accept contiguous
+            # replica groups for all-to-all -- which Ulysses needs -- while tolerating strided groups for
+            # all-reduce, which is all TP needs; since only one axis of a 2-D mesh can be contiguous, Ulysses
+            # gets it. On multi-node CUDA the opposite order is usually preferable, because TP is the most
+            # bandwidth-hungry collective and wants to stay within one NVLink domain. Pass `mesh=` on either
+            # config to choose the layout yourself.
+            mesh = (
+                cp_config.mesh
+                or tp_config.mesh
+                or torch.distributed.device_mesh.init_device_mesh(
+                    device_type=device_type,
+                    mesh_shape=(tp_config.tp_degree, *cp_config.mesh_shape),
+                    mesh_dim_names=("tp", *cp_config.mesh_dim_names),
+                )
+            )
+        elif cp_config is not None:
             mesh = cp_config.mesh or torch.distributed.device_mesh.init_device_mesh(
                 device_type=device_type,
                 mesh_shape=cp_config.mesh_shape,
                 mesh_dim_names=cp_config.mesh_dim_names,
             )
-        elif config.tensor_parallel_config is not None:
-            tp_config = config.tensor_parallel_config
+        elif tp_config is not None:
             mesh = tp_config.mesh or torch.distributed.device_mesh.init_device_mesh(
                 device_type=device_type,
                 mesh_shape=(tp_config.tp_degree,),
@@ -1736,6 +1757,32 @@ def _resolve_parallel_config(
 
         # `config.setup()` records the mesh resolved above onto the config; see `ParallelConfig.setup`.
         config.setup(rank, world_size, device, mesh=mesh)
+
+        # Validate the combination up front — after `setup`, which resolves `_tp_degree` from the mesh, but before
+        # anything is recorded on the model or applied to it. CP hooks are applied before TP, so a check left to the
+        # TP branch would raise on a model already carrying CP hooks: half-parallelised, and not usable.
+        if cp_config is not None and tp_config is not None:
+            tp_degree = tp_config._tp_degree
+            requested = cp_config.ring_degree * cp_config.ulysses_degree * tp_degree
+            if requested > world_size:
+                raise ValueError(
+                    f"Combining context and tensor parallelism needs `ring_degree` ({cp_config.ring_degree}) * "
+                    f"`ulysses_degree` ({cp_config.ulysses_degree}) * `tp_degree` ({tp_degree}) = {requested} "
+                    f"devices, which exceeds the world size ({world_size})."
+                )
+            num_heads = getattr(self.config, "num_attention_heads", None)
+            divisor = tp_degree * cp_config.ulysses_degree
+            if num_heads is not None and not cp_config.ulysses_anything and num_heads % divisor != 0:
+                # Ulysses trades sequence for heads inside attention (`SeqAllToAllDim` scatters over dim 2), and it
+                # only ever sees the heads TP left on this rank, so the count has to survive both splits.
+                raise ValueError(
+                    f"Combining tensor parallelism (`tp_degree`={tp_degree}) with Ulysses context parallelism "
+                    f"(`ulysses_degree`={cp_config.ulysses_degree}) requires the number of attention heads "
+                    f"({num_heads}) to be divisible by their product ({divisor}): TP shards the heads first, and "
+                    f"Ulysses splits what is left on each rank. Pass `ulysses_anything=True` to pad the head "
+                    f"dimension instead, or pick degrees whose product divides {num_heads}."
+                )
+
         self._parallel_config = config
         return config
 
@@ -1749,11 +1796,6 @@ def enable_parallelism(
             "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning."
         )
 
-        from ..hooks.context_parallel import apply_context_parallel
-        from .attention import AttentionModuleMixin
-        from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
-        from .attention_processor import Attention, MochiAttention
-
         if self._parallel_config is not None:
             raise RuntimeError(
                 f"Parallelism is already applied to this {self.__class__.__name__}. `enable_parallelism` cannot be "
@@ -1780,6 +1822,32 @@ def enable_parallelism(
 
         config = self._resolve_parallel_config(config)
 
+        self._apply_context_parallel(config, cp_plan)
+
+        if config.tensor_parallel_config is not None:
+            from ..hooks.tensor_parallel import apply_tensor_parallel
+
+            apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan)
+
+    def _apply_context_parallel(
+        self,
+        config: ParallelConfig,
+        cp_plan: dict[str, ContextParallelModelPlan] | None = None,
+    ):
+        """Install the context-parallel hooks `config` asks for, if any, and hand it to the attention processors.
+
+        Split out of `enable_parallelism` for the same reason as `_resolve_parallel_config`: `from_pretrained(...,
+        parallel_config=...)` shards the weights while reading them, so it cannot call `enable_parallelism`
+        afterwards, yet it still has to apply these hooks when the config asks for context parallelism.
+        """
+        from ..hooks.context_parallel import apply_context_parallel
+        from .attention import AttentionModuleMixin
+        from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
+        from .attention_processor import Attention, MochiAttention
+
+        if config.context_parallel_config is None:
+            return
+
         attention_classes = (Attention, MochiAttention, AttentionModuleMixin)
 
         if config.context_parallel_config is not None:
@@ -1829,11 +1897,6 @@ def enable_parallelism(
             cp_plan = cp_plan if cp_plan is not None else self._cp_plan
             apply_context_parallel(self, config.context_parallel_config, cp_plan)
 
-        if config.tensor_parallel_config is not None:
-            from ..hooks.tensor_parallel import apply_tensor_parallel
-
-            apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan)
-
     @classmethod
     def _load_pretrained_model(
         cls,