Allow tensor parallelism and context parallelism in one ParallelConfig

src/diffusers/hooks/tensor_parallel_neuron.py

@@ -22,7 +22,6 @@
 """
 
 import torch
-import torch.distributed as dist
 import torch.nn as nn
 
 
@@ -171,7 +170,11 @@ def _apply_tp_neuron(
 
     Model weights must be on CPU when this is called.
     """
-    rank = dist.get_rank()
+    # The shard index is this rank's coordinate on the tensor-parallel mesh, which is the global rank only
+    # when that mesh spans the whole world. Sharing a mesh with another parallelism makes `tp_mesh` a sub-mesh,
+    # and a global rank then indexes past the end of the weight. The generic backend already reads the
+    # coordinate (`tensor_parallel.py`), as does `ContextParallelConfig.setup`.
+    rank = tp_mesh.get_local_rank()
     tp_size = tp_mesh.size()
 
     for block, relative_plan in groups:

src/diffusers/models/_modeling_parallel.py

@@ -214,10 +214,10 @@ 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`."
             )
 
     def setup(
@@ -233,9 +233,14 @@ def setup(
         self._device = device
         self._mesh = mesh
         if self.context_parallel_config is not None:
+            # `ContextParallelConfig.setup` selects its own "ring" and "ulysses" dimensions out of `mesh`.
             self.context_parallel_config.setup(rank, world_size, device, mesh)
         if self.tensor_parallel_config is not None:
-            self.tensor_parallel_config.setup(rank, world_size, device, mesh)
+            # `TensorParallelConfig.setup` reads `tp_degree` off `mesh.size()`, so when the mesh is shared with
+            # context parallelism it must be handed the "tp" dimension alone rather than the whole mesh.
+            dim_names = (mesh.mesh_dim_names or ()) if mesh is not None else ()
+            tp_mesh = mesh["tp"] if "tp" in dim_names and len(dim_names) > 1 else mesh
+            self.tensor_parallel_config.setup(rank, world_size, device, tp_mesh)
 
 
 @dataclass(frozen=True)

src/diffusers/models/modeling_utils.py

@@ -1669,15 +1669,33 @@ def enable_parallelism(
                 break
 
         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,),

tests/models/testing_utils/__init__.py

@@ -20,6 +20,7 @@
 from .parallelism import (
     ContextParallelAttentionBackendsTesterMixin,
     ContextParallelTesterMixin,
+    HybridParallelTesterMixin,
     TensorParallelTesterMixin,
 )
 from .quantization import (
@@ -64,6 +65,7 @@
     "CacheTesterMixin",
     "ContextParallelTesterMixin",
     "ContextParallelAttentionBackendsTesterMixin",
+    "HybridParallelTesterMixin",
     "TensorParallelTesterMixin",
     "CPUOffloadTesterMixin",
     "FasterCacheConfigMixin",

tests/models/testing_utils/parallelism.py

@@ -21,7 +21,7 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from diffusers.models._modeling_parallel import ContextParallelConfig, TensorParallelConfig
+from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig, TensorParallelConfig
 from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
 
 from ...testing_utils import (
@@ -295,6 +295,63 @@ def _tensor_parallel_worker(
             dist.destroy_process_group()
 
 
+def _hybrid_parallel_worker(
+    rank, world_size, master_port, model_class, init_dict, cp_dict, tp_degree, inputs_dict, return_dict, state_dict
+):
+    """Worker function for combined tensor + context parallel inference testing.
+
+    Both parallelisms are requested through a single `ParallelConfig`, which shares one device mesh between them, so
+    each rank holds `1 / tp_degree` of every sharded weight *and* `1 / (ring_degree * ulysses_degree)` of the
+    sequence. Rank 0 reports its output so the caller can compare it against a single-device reference: the
+    composition is mathematically equivalent to the unsharded model up to floating-point reduction order.
+    """
+    try:
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(master_port)
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+
+        device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
+        backend = device_config["backend"]
+        device_module = device_config["module"]
+
+        dist.init_process_group(backend=backend, rank=rank, world_size=world_size)
+
+        device_module.set_device(rank)
+        device = torch.device(f"{torch_device}:{rank}")
+
+        model = model_class(**init_dict)
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+
+        inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+
+        model.enable_parallelism(
+            config=ParallelConfig(
+                tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree),
+                context_parallel_config=ContextParallelConfig(**cp_dict),
+            )
+        )
+
+        with torch.no_grad():
+            output = model(**inputs_on_device, return_dict=False)[0]
+
+        if rank == 0:
+            return_dict["status"] = "success"
+            return_dict["output_shape"] = list(output.shape)
+            # Serialise via nested list so the manager dict can transport it across processes.
+            return_dict["output"] = output.float().cpu().tolist()
+
+    except Exception as e:
+        if rank == 0:
+            return_dict["status"] = "error"
+            return_dict["error"] = str(e)
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
 @is_tensor_parallel
 @require_torch_multi_accelerator
 class TensorParallelTesterMixin:
@@ -345,6 +402,78 @@ def test_tensor_parallel_batch_inputs(self):
         self.test_tensor_parallel_inference(batch_size=2)
 
 
+@is_context_parallel
+@is_tensor_parallel
+@require_torch_multi_accelerator
+class HybridParallelTesterMixin:
+    """Tensor parallelism and context parallelism together, from one `ParallelConfig`.
+
+    Needs `tp_degree * ulysses_degree` accelerators (4 at the degrees used here), so it skips on a 2-device runner.
+    """
+
+    def test_hybrid_parallel_inference(self, batch_size: int = 1):
+        if not torch.distributed.is_available():
+            pytest.skip("torch.distributed is not available.")
+
+        for plan in ("_tp_plan", "_cp_plan"):
+            if getattr(self.model_class, plan, None) is None:
+                pytest.skip(f"Model does not define a `{plan}`, which hybrid parallelism requires.")
+
+        tp_degree, ulysses_degree = 2, 2
+        world_size = tp_degree * ulysses_degree
+        device_count = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])["module"].device_count()
+        if device_count < world_size:
+            pytest.skip(
+                f"tp_degree={tp_degree} x ulysses_degree={ulysses_degree} needs {world_size} accelerators, "
+                f"found {device_count}."
+            )
+
+        init_dict = self.get_init_dict()
+        num_heads = init_dict.get("num_attention_heads")
+        # Each rank keeps `num_heads // tp_degree` heads, which Ulysses splits again.
+        if num_heads is not None and num_heads % world_size != 0:
+            pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by {world_size}.")
+
+        inputs_dict = self.get_dummy_inputs(batch_size=batch_size)
+
+        # Single-device reference
+        model = self.model_class(**init_dict).eval().to(torch_device)
+        state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
+        with torch.no_grad():
+            ref_output = model(**inputs_dict, return_dict=False)[0].float().cpu()
+
+        inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+
+        master_port = _find_free_port()
+        manager = mp.Manager()
+        return_dict = manager.dict()
+
+        mp.spawn(
+            _hybrid_parallel_worker,
+            args=(
+                world_size,
+                master_port,
+                self.model_class,
+                init_dict,
+                {"ulysses_degree": ulysses_degree},
+                tp_degree,
+                inputs_dict,
+                return_dict,
+                state_dict,
+            ),
+            nprocs=world_size,
+            join=True,
+        )
+
+        assert return_dict.get("status") == "success", (
+            f"Hybrid parallel inference failed: {return_dict.get('error', 'Unknown error')}"
+        )
+
+        output = torch.tensor(return_dict["output"])
+        # Sharded matmuls plus the Ulysses all-to-all reorder the summation, hence the tolerance.
+        torch.testing.assert_close(ref_output, output, atol=1e-3, rtol=1e-3)
+
+
 @is_context_parallel
 @require_torch_multi_accelerator
 class ContextParallelTesterMixin:

tests/models/transformers/_neuron_hybrid_worker.py

@@ -0,0 +1,125 @@
+# coding=utf-8
+# Copyright 2026 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generic torchrun worker: assert a model's Neuron tensor-parallel x context-parallel output matches its reference.
+
+The counterpart of `_neuron_tp_worker.py` for the two parallelisms composed in a single `ParallelConfig`. Same
+contract: the model under test is supplied as a `module:function` spec reference on the command line, and the
+referenced factory returns `(model_class, init_dict, inputs)` with CPU tensors.
+
+    torchrun --nproc_per_node=8 _neuron_hybrid_worker.py \\
+        tests.models.transformers.test_models_transformer_flux:make_neuron_hybrid_spec
+
+`tp_degree` and `ulysses_degree` are read from `TP_DEGREE` / `ULYSSES_DEGREE` (defaults 2 and 4, whose product is
+the launched world size). `ulysses_degree` cannot be 2 on Neuron: its all-to-all only accepts group sizes of 4, 8,
+16 or multiples of 32.
+
+No mesh is passed, so this also exercises the default mesh layout that `enable_parallelism` builds for the combined
+case -- which matters on Neuron, where the all-to-all Ulysses depends on rejects strided replica groups and so the
+context-parallel dimensions have to vary fastest.
+
+Exit code 0 means the composed path is numerically equivalent to the unsharded model; non-zero means failure.
+"""
+
+import argparse
+import importlib
+import os
+import sys
+import traceback
+
+
+# Make the in-repo `diffusers` and `tests` packages importable when run via torchrun from an arbitrary CWD.
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "src"))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+
+import torch
+import torch.distributed as dist
+import torch_neuronx  # noqa: F401 — registers torch.neuron
+
+from diffusers import ContextParallelConfig, ParallelConfig, TensorParallelConfig
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Neuron tensor-parallel x context-parallel correctness worker.")
+    parser.add_argument(
+        "spec",
+        help="`module:function` reference returning (model_class, init_dict, cpu_inputs) for the model under test.",
+    )
+    args = parser.parse_args()
+    module_name, _, fn_name = args.spec.partition(":")
+    model_class, init_dict, inputs = getattr(importlib.import_module(module_name), fn_name)()
+
+    tp_degree = int(os.environ.get("TP_DEGREE", "2"))
+    ulysses_degree = int(os.environ.get("ULYSSES_DEGREE", "4"))
+
+    dist.init_process_group(backend="neuron")
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    device = torch.neuron.current_device()
+
+    if tp_degree * ulysses_degree != world_size:
+        raise ValueError(
+            f"tp_degree ({tp_degree}) x ulysses_degree ({ulysses_degree}) must equal the world size ({world_size})."
+        )
+
+    # Identical weights on every rank (same seed), kept on CPU as the Neuron pre-shard backend requires.
+    torch.manual_seed(0)
+    model = model_class(**init_dict).eval()
+
+    # Single-device (unsharded) reference on CPU, computed before the shard plan mutates the weights in place.
+    with torch.no_grad():
+        ref_output = model(**inputs, return_dict=False)[0].float().cpu()
+
+    model.enable_parallelism(
+        config=ParallelConfig(
+            tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree),
+            context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree),
+        )
+    )
+    model = model.to(device)
+    torch.neuron.synchronize()
+
+    inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+    with torch.no_grad():
+        output = model(**inputs_on_device, return_dict=False)[0]
+    torch.neuron.synchronize()
+    output = output.float().cpu()
+
+    if rank == 0:
+        assert output.shape == ref_output.shape, f"shape mismatch: {output.shape} vs {ref_output.shape}"
+        assert torch.isfinite(output).all(), "output contains non-finite values"
+        max_abs = (output - ref_output).abs().max().item()
+        denom = ref_output.abs().max().item() + 1e-6
+        print(
+            f"[rank0] tp_degree={tp_degree} ulysses_degree={ulysses_degree} "
+            f"output_shape={tuple(output.shape)} max_abs_diff={max_abs:.4e} max_rel_diff={max_abs / denom:.4e}"
+        )
+        # Neuron runs matmuls in bf16 internally, so compare with a bf16-level tolerance, as `_neuron_tp_worker`
+        # does. A wrong shard plan or a mis-ordered mesh produces grossly different output and is caught well
+        # inside this bound.
+        torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2)
+        print("[rank0] PASS: Neuron hybrid-parallel output matches single-device reference.")
+
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception:
+        traceback.print_exc()
+        # Ensure a non-zero exit so the launching pytest sees the failure.
+        os._exit(1)