diff --git a/src/diffusers/hooks/tensor_parallel_neuron.py b/src/diffusers/hooks/tensor_parallel_neuron.py index 6b8f219a17ff..9c53afc0f623 100644 --- a/src/diffusers/hooks/tensor_parallel_neuron.py +++ b/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: diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index 86627284e078..a98574cc49e8 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/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) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 425f2f29235e..6926b6fa89d3 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/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,), diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 760a3fac04e0..4d70f3a1ae29 100644 --- a/tests/models/testing_utils/__init__.py +++ b/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", diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index 63575abf6b7b..49c7acda6531 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 ...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: 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/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 53af9eedc50c..e4d7ce8b808e 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, @@ -268,6 +275,23 @@ class TestFluxTransformerTensorParallel(FluxTransformerTesterConfig, TensorParal """Tensor Parallel inference tests for Flux Transformer (CUDA/XPU multi-accelerator).""" +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`). @@ -301,6 +325,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."""