diff --git a/src/diffusers/hooks/context_parallel.py b/src/diffusers/hooks/context_parallel.py index 1310b20c5c11..4b0a90dac4c2 100644 --- a/src/diffusers/hooks/context_parallel.py +++ b/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] diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index 58cbc822f8bf..cf576ab6bb77 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/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) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index a6881e971edf..e6c1ef6f8266 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/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) @@ -1720,15 +1723,33 @@ def _resolve_parallel_config( device = torch.device(device_type, rank % device_module.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,), @@ -1737,6 +1758,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 @@ -1750,11 +1797,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 " @@ -1781,6 +1823,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: @@ -1830,11 +1898,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, diff --git a/tests/models/test_parallelism_combined.py b/tests/models/test_parallelism_combined.py new file mode 100644 index 000000000000..4a404a06125e --- /dev/null +++ b/tests/models/test_parallelism_combined.py @@ -0,0 +1,164 @@ +# 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. + +"""Tests for combining context parallelism with tensor parallelism on one device mesh. + +These cover the wiring rather than the numerics: that `ParallelConfig` hands each parallelism its own submesh, and +that the resulting process groups are the intended factorisation of the world. They run on CPU over `gloo`, so they +need no accelerator — `gloo` has no `all_to_all`, so an end-to-end Ulysses forward pass cannot run here; that is +covered by `ContextAndTensorParallelTesterMixin` in `testing_utils/parallelism.py`, which needs four accelerators. +""" + +import os +import socket + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig, TensorParallelConfig + + +def _find_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.listen(1) + return s.getsockname()[1] + + +def _mesh_factorization_worker(rank, world_size, master_port, ulysses_degree, tp_degree, return_dict): + """Build a combined mesh, run `ParallelConfig.setup`, and report the process groups each parallelism landed on.""" + 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) + + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + mesh = torch.distributed.device_mesh.init_device_mesh( + "cpu", mesh_shape=(tp_degree, 1, ulysses_degree), mesh_dim_names=("tp", "ring", "ulysses") + ) + config = ParallelConfig( + context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree), + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + ) + config.setup(rank, world_size, torch.device("cpu"), mesh=mesh) + + cp_config = config.context_parallel_config + tp_config = config.tensor_parallel_config + return_dict[rank] = { + "status": "success", + "cp_ranks": dist.get_process_group_ranks(cp_config._flattened_mesh.get_group()), + "ulysses_local_rank": cp_config._ulysses_local_rank, + "tp_ranks": dist.get_process_group_ranks(tp_config._mesh.get_group()), + "tp_local_rank": tp_config._mesh.get_local_rank(), + "tp_degree": tp_config._tp_degree, + } + except Exception as e: # noqa: BLE001 — surfaced as a test failure by the caller + return_dict[rank] = {"status": "error", "error": f"{type(e).__name__}: {e}"} + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _missing_tp_dim_worker(rank, world_size, master_port, return_dict): + """A combined config handed a CP-only mesh must be rejected, naming the missing 'tp' dimension.""" + 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) + + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + mesh = torch.distributed.device_mesh.init_device_mesh( + "cpu", mesh_shape=(1, world_size), mesh_dim_names=("ring", "ulysses") + ) + config = ParallelConfig( + context_parallel_config=ContextParallelConfig(ulysses_degree=world_size), + tensor_parallel_config=TensorParallelConfig(tp_degree=1), + ) + try: + config.setup(rank, world_size, torch.device("cpu"), mesh=mesh) + except ValueError as e: + return_dict[rank] = {"status": "raised", "message": str(e)} + else: + return_dict[rank] = {"status": "no_raise"} + except Exception as e: # noqa: BLE001 + return_dict[rank] = {"status": "error", "error": f"{type(e).__name__}: {e}"} + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _spawn(worker, world_size, *args): + if not dist.is_available(): + pytest.skip("torch.distributed is not available.") + manager = mp.Manager() + return_dict = manager.dict() + mp.spawn(worker, args=(world_size, _find_free_port(), *args, return_dict), nprocs=world_size, join=True) + return return_dict + + +class TestCombinedParallelConfig: + """`ParallelConfig` accepting both parallelisms at once, and splitting the mesh between them.""" + + def test_both_configs_are_accepted(self): + # Combining the two used to raise outright; the mesh dimension per parallelism is what makes it work. + config = ParallelConfig( + context_parallel_config=ContextParallelConfig(ulysses_degree=2), + tensor_parallel_config=TensorParallelConfig(tp_degree=2), + ) + assert config._is_combined + assert config.context_parallel_config.ulysses_degree == 2 + assert config.tensor_parallel_config.tp_degree == 2 + + def test_single_parallelism_is_not_combined(self): + assert not ParallelConfig(context_parallel_config=ContextParallelConfig(ulysses_degree=2))._is_combined + assert not ParallelConfig(tensor_parallel_config=TensorParallelConfig(tp_degree=2))._is_combined + + def test_mesh_is_factored_between_the_two(self): + """Each parallelism must end up on its own process group: TP within a CP chunk, CP across the TP groups. + + With the default `mesh_shape=(tp=2, ring=1, ulysses=2)` over four ranks, "ulysses" is the fastest-varying + dimension, so ranks {0,1} and {2,3} are the CP groups and {0,2} / {1,3} are the TP groups. A TP all-reduce + must not reach a rank holding a different sequence chunk, and a CP collective must not reach a rank holding a + different weight shard — this asserts exactly that split. + """ + world_size = 4 + results = _spawn(_mesh_factorization_worker, world_size, 2, 2) + + for rank in range(world_size): + assert results[rank]["status"] == "success", results[rank].get("error") + + assert [results[r]["tp_ranks"] for r in range(4)] == [[0, 2], [1, 3], [0, 2], [1, 3]] + assert [results[r]["cp_ranks"] for r in range(4)] == [[0, 1], [0, 1], [2, 3], [2, 3]] + assert [results[r]["ulysses_local_rank"] for r in range(4)] == [0, 1, 0, 1] + # The TP shard index is the coordinate inside the TP group, which stops tracking the global rank as soon as + # the mesh has more than one dimension. Weight sharding keys off this, so it is the value that matters. + assert [results[r]["tp_local_rank"] for r in range(4)] == [0, 0, 1, 1] + assert all(results[r]["tp_degree"] == 2 for r in range(4)) + + def test_mesh_without_tp_dimension_is_rejected(self): + world_size = 2 + results = _spawn(_missing_tp_dim_worker, world_size) + + for rank in range(world_size): + assert results[rank]["status"] == "raised", ( + f"rank {rank}: expected a ValueError for a mesh with no 'tp' dimension, got {results[rank]}" + ) + assert "tp" in results[rank]["message"] diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index ee0ab21e8719..6db810342821 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -21,8 +21,10 @@ from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin from .parallelism import ( + ContextAndTensorParallelTesterMixin, ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, + HybridParallelTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, ) @@ -66,8 +68,10 @@ "BitsAndBytesConfigMixin", "BitsAndBytesTesterMixin", "CacheTesterMixin", + "ContextAndTensorParallelTesterMixin", "ContextParallelTesterMixin", "ContextParallelAttentionBackendsTesterMixin", + "HybridParallelTesterMixin", "TensorParallelTesterMixin", "TensorParallelTPUTesterMixin", "CPUOffloadTesterMixin", diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index eca6bff5d271..d7f13da16cce 100644 --- a/tests/models/testing_utils/parallelism.py +++ b/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 diffusers.utils import SAFE_WEIGHTS_INDEX_NAME @@ -380,6 +380,63 @@ def _tensor_parallel_from_pretrained_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: @@ -558,6 +615,222 @@ def test_tensor_parallel_tpu_inference(self, atol=1e-3, rtol=1e-3): torch.testing.assert_close(ref_output, tp_output, atol=atol, rtol=rtol) +def _context_and_tensor_parallel_worker( + rank, world_size, master_port, model_class, init_dict, cp_dict, tp_degree, inputs_dict, return_dict, state_dict +): + """Worker for combined context + tensor parallel inference. + + Both configs go into one `ParallelConfig`, which puts them on one mesh with a dimension each. The result should + still match the single-device reference: TP is mathematically equivalent to the unsharded model, and CP splits the + sequence and gathers it back, so neither changes the function being computed. + """ + 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"]) + dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size) + + device_config["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( + context_parallel_config=ContextParallelConfig(**cp_dict), + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + ) + ) + + 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) + return_dict["output"] = output.float().cpu().tolist() + + except Exception as e: + if rank == 0: + return_dict["status"] = "error" + return_dict["error"] = f"{type(e).__name__}: {e}" + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +@is_context_parallel +@is_tensor_parallel +@require_torch_multi_accelerator +class ContextAndTensorParallelTesterMixin: + """Context and tensor parallelism enabled together, over one device mesh. + + The two are orthogonal — CP cuts the sequence, TP cuts the weights — so composing them should leave the computed + function unchanged. Both axes are covered: Ulysses, which trades sequence for heads inside attention, and ring, + which does not touch the head dimension at all. + + Needs `cp_degree` x `tp_degree` accelerators (four by default). Head-count requirements differ per CP type, so a + model whose dummy config has too few heads skips rather than failing. + """ + + cp_degree = 2 + tp_degree = 2 + + @pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"]) + def test_context_and_tensor_parallel_inference(self, cp_type, batch_size: int = 1): + if not torch.distributed.is_available(): + pytest.skip("torch.distributed is not available.") + + if getattr(self.model_class, "_tp_plan", None) is None: + pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.") + if getattr(self.model_class, "_cp_plan", None) is None: + pytest.skip("Model does not define a `_cp_plan` for context parallel inference.") + + if cp_type == "ring_degree": + active_backend, _ = _AttentionBackendRegistry.get_active_backend() + if active_backend == AttentionBackendName.NATIVE: + pytest.skip("Ring attention is not supported with the native attention backend.") + + world_size = self.cp_degree * self.tp_degree + device_module = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])["module"] + if device_module.device_count() < world_size: + pytest.skip(f"Combined CP x TP needs {world_size} accelerators, found {device_module.device_count()}.") + + init_dict = self.get_init_dict() + num_heads = init_dict.get("num_attention_heads") + # TP shards the heads; Ulysses then splits what TP left on each rank, so the two multiply. Ring leaves the + # head dimension alone, so only the TP degree has to divide the head count. + required_head_multiple = self.tp_degree * (self.cp_degree if cp_type == "ulysses_degree" else 1) + if num_heads is not None and num_heads % required_head_multiple != 0: + pytest.skip( + f"`num_attention_heads` ({num_heads}) is not divisible by {required_head_multiple}, required for " + f"{cp_type.removesuffix('_degree')}={self.cp_degree} x tp_degree={self.tp_degree}." + ) + + inputs_dict = self.get_dummy_inputs(batch_size=batch_size) + + # Single-device reference, captured before anything is sharded. + 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()} + + manager = mp.Manager() + return_dict = manager.dict() + mp.spawn( + _context_and_tensor_parallel_worker, + args=( + world_size, + _find_free_port(), + self.model_class, + init_dict, + {cp_type: self.cp_degree}, + self.tp_degree, + inputs_dict, + return_dict, + state_dict, + ), + nprocs=world_size, + join=True, + ) + + assert return_dict.get("status") == "success", ( + f"Combined context + tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}" + ) + + combined_output = torch.tensor(return_dict["output"]) + assert list(ref_output.shape) == return_dict["output_shape"] + # Two sets of collectives reorder the summation on top of the sharded matmuls, so the tolerance matches the + # TP-only test rather than the tighter CP-only one. + torch.testing.assert_close(ref_output, combined_output, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"], ids=["ulysses", "ring"]) + def test_context_and_tensor_parallel_batch_inputs(self, cp_type): + self.test_context_and_tensor_parallel_inference(cp_type, 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: diff --git a/tests/models/transformers/_neuron_hybrid_worker.py b/tests/models/transformers/_neuron_hybrid_worker.py new file mode 100644 index 000000000000..1c054af44104 --- /dev/null +++ b/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) diff --git a/tests/models/transformers/_neuron_sharded_load_worker.py b/tests/models/transformers/_neuron_sharded_load_worker.py new file mode 100644 index 000000000000..dd25d33b35a9 --- /dev/null +++ b/tests/models/transformers/_neuron_sharded_load_worker.py @@ -0,0 +1,170 @@ +# 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 `from_pretrained(..., parallel_config=...)` installs the context-parallel hooks. + +The sharded-load path shards weights while reading the checkpoint, so it cannot call `enable_parallelism` afterwards +and applies the parallelism itself. Tensor parallelism it did apply; the context-parallel hooks it did not, which is +silent — the model still returns the right numbers, every rank just redundantly computes the whole sequence. So this +worker checks both halves: the hooks are registered *and* the output still matches a single-device reference. + +Same contract as `_neuron_tp_worker.py`: 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_sharded_load_worker.py \\ + tests.models.transformers.test_models_transformer_flux:make_neuron_sharded_load_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. + +Rank 0 writes the checkpoint that every rank then reads, so this is single-node only, as the rest of the Neuron +test workers are. + +Exit code 0 means the sharded-load path applied both parallelisms correctly; non-zero means failure. +""" + +import argparse +import importlib +import os +import shutil +import sys +import tempfile +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 _context_parallel_hook_names(model): + """Every `cp_input---*` / `cp_output---*` hook `apply_context_parallel` registered on `model`.""" + names = [] + for module in model.modules(): + registry = getattr(module, "_diffusers_hook", None) + if registry is not None: + names.extend(name for name in registry.hooks if name.startswith(("cp_input", "cp_output"))) + return names + + +def main(): + parser = argparse.ArgumentParser(description="Neuron sharded-load context-parallel hook 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})." + ) + + # Rank 0 writes the checkpoint the sharded loader then reads on every rank. Safetensors, because that is what + # the per-rank streaming loader requires. + checkpoint_dir = os.environ.get( + "SHARDED_LOAD_CHECKPOINT_DIR", os.path.join(tempfile.gettempdir(), "diffusers_neuron_sharded_load_ckpt") + ) + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + torch.manual_seed(0) + model_class(**init_dict).eval().save_pretrained(checkpoint_dir) + dist.barrier() + + # Single-device (unsharded) reference, read back from the same checkpoint so only the loading path differs. + ref_output = None + if rank == 0: + reference = model_class.from_pretrained(checkpoint_dir).eval() + with torch.no_grad(): + ref_output = reference(**inputs, return_dict=False)[0].float().cpu() + del reference + + model = model_class.from_pretrained( + checkpoint_dir, + parallel_config=ParallelConfig( + tensor_parallel_config=TensorParallelConfig(tp_degree=tp_degree), + context_parallel_config=ContextParallelConfig(ulysses_degree=ulysses_degree), + ), + ).eval() + torch.neuron.synchronize() + + # The point of the test: the loader has to install these itself, since `enable_parallelism` is unavailable + # once the weights are sharded. + hook_names = _context_parallel_hook_names(model) + assert hook_names, ( + "`from_pretrained(..., parallel_config=...)` applied tensor parallelism but registered no context-parallel " + "hooks, so the `context_parallel_config` was silently ignored." + ) + processors = [getattr(module, "processor", None) for module in model.modules()] + processor_configs = [ + getattr(p, "_parallel_config", None) for p in processors if p is not None and hasattr(p, "_parallel_config") + ] + assert processor_configs and all(c is not None for c in processor_configs), ( + "Context-parallel hooks are registered but the attention processors did not receive the `ParallelConfig`, " + "so attention would run without the Ulysses all-to-all." + ) + + 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"context_parallel_hooks={len(hook_names)} output_shape={tuple(output.shape)} " + f"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: sharded load applied both parallelisms and matches the single-device reference.") + + dist.barrier() + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + 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) diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 04542253c7f6..7942be68d323 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -27,7 +27,13 @@ from diffusers.models.transformers.transformer_flux import FluxIPAdapterAttnProcessor from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device +from ...testing_utils import ( + enable_full_determinism, + is_context_parallel, + is_tensor_parallel, + require_torch_neuron, + torch_device, +) from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -40,6 +46,7 @@ FirstBlockCacheTesterMixin, GGUFCompileTesterMixin, GGUFTesterMixin, + HybridParallelTesterMixin, IPAdapterTesterMixin, LoraHotSwappingForModelTesterMixin, LoraTesterMixin, @@ -285,6 +292,23 @@ def get_init_dict(self): return {**super().get_init_dict(), "num_attention_heads": self.tp_world_size} +class TestFluxTransformerHybridParallel(FluxTransformerTesterConfig, HybridParallelTesterMixin): + """Tensor Parallel x Context Parallel inference tests for Flux Transformer (needs 4 accelerators).""" + + +def make_neuron_hybrid_spec(): + """Model spec consumed by the generic Neuron hybrid worker (`_neuron_hybrid_worker.py`). + + Same contract as `make_neuron_tp_spec`, but `num_attention_heads` is raised to 8 so the head count survives + being divided twice: `tp_degree=2` leaves 4 heads per rank and `ulysses_degree=4` splits those into 1 each. + (`ulysses_degree` cannot be 2 on Neuron, whose all-to-all only accepts group sizes of 4, 8, 16 or multiples + of 32.) + """ + config = FluxTransformerTesterConfig() + init_dict = config.get_init_dict() | {"num_attention_heads": 8} + return FluxTransformer2DModel, init_dict, config.get_dummy_inputs(device="cpu") + + def make_neuron_tp_spec(): """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`). @@ -296,6 +320,41 @@ def make_neuron_tp_spec(): return FluxTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") +def make_neuron_sharded_load_spec(): + """Model spec consumed by the sharded-load worker (`_neuron_sharded_load_worker.py`). + + Same contract as `make_neuron_tp_spec`, but `num_attention_heads` is raised to 8 so the head count survives + being divided twice: `tp_degree=2` leaves 4 heads per rank and `ulysses_degree=4` splits those into 1 each. + (`ulysses_degree` cannot be 2 on Neuron, whose all-to-all only accepts group sizes of 4, 8, 16 or multiples + of 32.) + """ + config = FluxTransformerTesterConfig() + init_dict = config.get_init_dict() | {"num_attention_heads": 8} + return FluxTransformer2DModel, init_dict, config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_neuron +class TestFluxTransformerShardedLoadNeuron: + """`from_pretrained(..., parallel_config=...)` with both parallelisms, on AWS Neuron. + + Launched the same way as `TestFluxTransformerTensorParallelNeuron`, at `tp_degree=2 x ulysses_degree=4`, i.e. 8 + ranks. The worker asserts that the sharded-load path registered the context-parallel hooks — which numbers alone + cannot detect, since a model that silently skips them still returns the right answer — and that the output still + matches a single-device reference read back from the same checkpoint. + """ + + def test_sharded_load_context_parallel_neuron(self): + worker = os.path.join(os.path.dirname(__file__), "_neuron_sharded_load_worker.py") + spec = "tests.models.transformers.test_models_transformer_flux:make_neuron_sharded_load_spec" + cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=8", worker, spec] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, ( + f"Neuron sharded-load worker failed (exit {result.returncode}).\n" + f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" + ) + + @is_tensor_parallel @require_torch_neuron class TestFluxTransformerTensorParallelNeuron: @@ -318,6 +377,28 @@ def test_tensor_parallel_neuron_inference(self): ) +@is_context_parallel +@is_tensor_parallel +@require_torch_neuron +class TestFluxTransformerHybridParallelNeuron: + """Tensor Parallel x Context Parallel inference test for Flux Transformer on AWS Neuron. + + Same launching pattern as `TestFluxTransformerTensorParallelNeuron`: Neuron needs `torchrun` with the + `"neuron"` distributed backend, so it cannot use the `torch.multiprocessing`/NCCL spawn path of + `HybridParallelTesterMixin`. Runs at `tp_degree=2 x ulysses_degree=4`, i.e. 8 ranks. + """ + + def test_hybrid_parallel_neuron_inference(self): + worker = os.path.join(os.path.dirname(__file__), "_neuron_hybrid_worker.py") + spec = "tests.models.transformers.test_models_transformer_flux:make_neuron_hybrid_spec" + cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=8", worker, spec] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, ( + f"Neuron hybrid-parallel worker failed (exit {result.returncode}).\n" + f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" + ) + + class TestFluxTransformerIPAdapter(FluxTransformerTesterConfig, IPAdapterTesterMixin): """IP Adapter tests for Flux Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 6210aecff87f..10d36c8155ba 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -30,6 +30,7 @@ AttentionTesterMixin, BaseModelTesterConfig, BitsAndBytesTesterMixin, + ContextAndTensorParallelTesterMixin, ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, LoraHotSwappingForModelTesterMixin, @@ -320,6 +321,25 @@ def test_tensor_parallel_tpu_inference(self): super().test_tensor_parallel_tpu_inference(atol=2e-2, rtol=2e-2) +class TestQwenImageTransformerContextAndTensorParallel( + QwenImageTransformerTesterConfig, ContextAndTensorParallelTesterMixin +): + """Context Parallel x Tensor Parallel inference tests for QwenImage Transformer (4 accelerators). + + QwenImage is the model this runs on because its dummy config has four attention heads, and Ulysses splits the + heads TP already sharded — so `ulysses_degree` x `tp_degree` has to divide the head count. + """ + + def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: + inputs = super().get_dummy_inputs(batch_size=batch_size) + encoder_hidden_states_mask = inputs["encoder_hidden_states_mask"] + encoder_hidden_states_mask[:, 1] = 0 + encoder_hidden_states_mask[:, 3] = 0 + encoder_hidden_states_mask[:, 5:] = 0 + inputs["encoder_hidden_states_mask"] = encoder_hidden_states_mask + return inputs + + def make_neuron_tp_spec(): """Model spec consumed by the generic Neuron TP worker (``_neuron_tp_worker.py``).