diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index a6881e971edf..469906613a5e 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1716,8 +1716,7 @@ def _resolve_parallel_config( rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() device_type = torch._C._get_accelerator().type - device_module = torch.get_device_module(device_type) - device = torch.device(device_type, rank % device_module.device_count()) + device = torch.device(device_type, rank % torch.accelerator.device_count()) mesh = None if config.context_parallel_config is not None: diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index b6eabfd584bb..1a9ef250c519 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -337,9 +337,12 @@ def _qwenimage21_prepare_qkv( key = attn.to_k(hidden_states) value = attn.to_v(hidden_states) - query = query.unflatten(-1, (attn.heads, -1)) - key = key.unflatten(-1, (attn.heads, -1)) - value = value.unflatten(-1, (attn.heads, -1)) + # Split by `head_dim` rather than by `attn.heads`: under tensor parallelism each rank holds only its share of the + # heads, while `attn.heads` and `attn.inner_dim` keep their full values. + head_dim = attn.inner_dim // attn.heads + query = query.unflatten(-1, (-1, head_dim)) + key = key.unflatten(-1, (-1, head_dim)) + value = value.unflatten(-1, (-1, head_dim)) query = attn.norm_q(query).to(value.dtype) key = attn.norm_k(key).to(value.dtype) @@ -758,6 +761,18 @@ class QwenImage21Transformer2DModel( _skip_layerwise_casting_patterns = ["pos_embed", "norm"] _repeated_blocks = ["QwenImage21TransformerBlock"] _skip_keys = ["kv_cache"] + # Tensor-parallel plan: every block's attention and SwiGLU projections are separate, bias-free Linears, so each + # entry is a plain "colwise"/"rowwise" pair and no packed sharding is needed. The shared `modulation`, the input + # projections (`img_in`, `txt_in`), `norm_out` and `proj_out` stay replicated (intentionally absent here). + _tp_plan = { + "transformer_blocks.*.attn.to_q": "colwise", + "transformer_blocks.*.attn.to_k": "colwise", + "transformer_blocks.*.attn.to_v": "colwise", + "transformer_blocks.*.attn.to_out.0": "rowwise", + "transformer_blocks.*.img_mlp.proj": "colwise", + "transformer_blocks.*.img_mlp.gate_layer": "colwise", + "transformer_blocks.*.img_mlp.out": "rowwise", + } @register_to_config def __init__( @@ -833,9 +848,10 @@ def build_token_metadata( ) image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long) - block_ids = torch.repeat_interleave( - torch.arange(len(block_lengths), device=image_pad_mask.device), - torch.tensor(block_lengths, device=image_pad_mask.device), + # Neuron cannot lower a tensor-repeats repeat_interleave, so build the ids from the Python block lengths. + block_ids = torch.tensor( + [block for block, length in enumerate(block_lengths) for _ in range(length)], + device=image_pad_mask.device, ) image_ids[image_positions] = block_ids diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index ee0ab21e8719..a630a2c8b2c0 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -23,6 +23,7 @@ from .parallelism import ( ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, + TensorParallelNeuronTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, ) @@ -69,6 +70,7 @@ "ContextParallelTesterMixin", "ContextParallelAttentionBackendsTesterMixin", "TensorParallelTesterMixin", + "TensorParallelNeuronTesterMixin", "TensorParallelTPUTesterMixin", "CPUOffloadTesterMixin", "FasterCacheConfigMixin", diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index eca6bff5d271..de19a7edd212 100644 --- a/tests/models/testing_utils/parallelism.py +++ b/tests/models/testing_utils/parallelism.py @@ -31,6 +31,7 @@ is_kernels_available, is_tensor_parallel, require_torch_multi_accelerator, + require_torch_neuron, require_torch_tpu, torch_device, ) @@ -43,6 +44,7 @@ "cuda": {"backend": "nccl", "module": torch.cuda}, "xpu": {"backend": "xccl", "module": torch.xpu}, "tpu": {"backend": "tpu_dist", "module": None}, + "neuron": {"backend": "neuron", "module": getattr(torch, "neuron", None)}, } @@ -558,6 +560,59 @@ 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) +@is_tensor_parallel +@require_torch_neuron +class TensorParallelNeuronTesterMixin: + """Same check as `TensorParallelTesterMixin`, spawning one process per Neuron core. + + Run these tests in their own pytest process (e.g. `-k TensorParallelNeuron`): an earlier test that allocates a + Neuron tensor initializes the Neuron runtime in the pytest process, and the spawned workers then fail to initialize + theirs. + """ + + tp_world_size = 2 + + def test_tensor_parallel_neuron_inference(self, atol=2e-2, rtol=2e-2): + if getattr(self.model_class, "_tp_plan", None) is None: + pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.") + + world_size = self.tp_world_size + init_dict = self.get_init_dict() + num_heads = init_dict.get("num_attention_heads") + if num_heads is not None and num_heads % world_size != 0: + pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).") + + # Reference on CPU: touching the Neuron device here would break the spawned workers. + inputs_dict = self.get_dummy_inputs(device="cpu") + model = self.model_class(**init_dict).eval() + with torch.no_grad(): + ref_output = model(**inputs_dict, return_dict=False)[0].float() + + return_dict = mp.Manager().dict() + mp.spawn( + _tensor_parallel_worker, + args=( + world_size, + _find_free_port(), + self.model_class, + init_dict, + inputs_dict, + return_dict, + model.state_dict(), + "neuron", + ), + nprocs=world_size, + join=True, + ) + + assert return_dict.get("status") == "success", ( + f"Tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}" + ) + tp_output = torch.tensor(return_dict["output"]) + # Neuron runs matmuls in bf16 internally, so compare with a bf16-level tolerance. + torch.testing.assert_close(ref_output, tp_output, atol=atol, rtol=rtol) + + @is_context_parallel @require_torch_multi_accelerator class ContextParallelTesterMixin: diff --git a/tests/models/transformers/_neuron_tp_worker.py b/tests/models/transformers/_neuron_tp_worker.py deleted file mode 100644 index 681f5686bee9..000000000000 --- a/tests/models/transformers/_neuron_tp_worker.py +++ /dev/null @@ -1,111 +0,0 @@ -# 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 output matches its single-device reference. - -Model-agnostic. The model under test is supplied as a `module:function` spec reference on the command line; the -referenced factory returns `(model_class, init_dict, inputs)` with CPU tensors, so all model-specific test data lives -with the launching test rather than here. - -Launched as a subprocess by a `@require_torch_neuron` test (and runnable directly for debugging):: - - torchrun --nproc_per_node=2 _neuron_tp_worker.py \\ - tests.models.transformers.test_models_transformer_flux2:make_neuron_tp_spec - -Each rank builds an identical (seeded) model on CPU, computes a single-device reference, then shards it with -`enable_parallelism(TensorParallelConfig(mesh=neuron_mesh))` — which auto-selects the Neuron pre-shard backend — runs -a forward pass on the Neuron device, and asserts the gathered output matches the reference. Exit code 0 means the TP -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 torch.distributed.device_mesh import DeviceMesh - -from diffusers import TensorParallelConfig - - -def main(): - parser = argparse.ArgumentParser(description="Neuron tensor-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)() - - dist.init_process_group(backend="neuron") - rank = dist.get_rank() - tp_size = dist.get_world_size() - device = torch.neuron.current_device() - tp_mesh = DeviceMesh("neuron", list(range(tp_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 TP mutates the weights in place. - with torch.no_grad(): - ref_output = model(**inputs, return_dict=False)[0].float().cpu() - - # Shard across all ranks; the Neuron backend is auto-selected from the mesh device type. - model.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh)) - 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(): - tp_output = model(**inputs_on_device, return_dict=False)[0] - torch.neuron.synchronize() - tp_output = tp_output.float().cpu() - - if rank == 0: - assert tp_output.shape == ref_output.shape, f"shape mismatch: {tp_output.shape} vs {ref_output.shape}" - assert torch.isfinite(tp_output).all(), "TP output contains non-finite values" - max_abs = (tp_output - ref_output).abs().max().item() - denom = ref_output.abs().max().item() + 1e-6 - print( - f"[rank0] tp_size={tp_size} output_shape={tuple(tp_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. A wrong shard - # plan produces grossly different output and is caught comfortably within this bound. - torch.testing.assert_close(tp_output, ref_output, atol=2e-2, rtol=2e-2) - print("[rank0] PASS: Neuron tensor-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 04542253c7f6..827987cf47c8 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -import subprocess -import sys import tempfile from typing import Any @@ -27,7 +24,7 @@ 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, torch_device from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -52,6 +49,7 @@ SDNQTesterMixin, SingleFileTesterMixin, TaylorSeerCacheTesterMixin, + TensorParallelNeuronTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, TorchAoCompileTesterMixin, @@ -285,37 +283,8 @@ def get_init_dict(self): return {**super().get_init_dict(), "num_attention_heads": self.tp_world_size} -def make_neuron_tp_spec(): - """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`). - - Returns `(model_class, init_dict, cpu_inputs)`. Defined here so all Flux-specific test data lives in this file - while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest of - the Flux tests. - """ - config = FluxTransformerTesterConfig() - return FluxTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") - - -@is_tensor_parallel -@require_torch_neuron -class TestFluxTransformerTensorParallelNeuron: - """Tensor Parallel inference test for Flux Transformer on AWS Neuron. - - Neuron TP runs through `torchrun` with the `"neuron"` distributed backend, so it cannot use the - `torch.multiprocessing`/NCCL spawn path of `TensorParallelTesterMixin`. This launches the generic worker - with the Flux model spec (`make_neuron_tp_spec`); the worker asserts the sharded output matches a single-device - reference, and the test checks its exit code. - """ - - def test_tensor_parallel_neuron_inference(self): - worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py") - spec = "tests.models.transformers.test_models_transformer_flux:make_neuron_tp_spec" - cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec] - result = subprocess.run(cmd, capture_output=True, text=True) - assert result.returncode == 0, ( - f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n" - f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" - ) +class TestFluxTransformerTensorParallelNeuron(FluxTransformerTesterConfig, TensorParallelNeuronTesterMixin): + """Tensor Parallel inference test for Flux Transformer on AWS Neuron.""" class TestFluxTransformerIPAdapter(FluxTransformerTesterConfig, IPAdapterTesterMixin): diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index ccfea457606e..966a3c746bd6 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -13,10 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -import subprocess -import sys - import pytest import torch @@ -30,7 +26,7 @@ ) 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, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -45,6 +41,7 @@ MemoryTesterMixin, ModelTesterMixin, SingleFileTesterMixin, + TensorParallelNeuronTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, TorchAoCompileTesterMixin, @@ -171,17 +168,6 @@ class TestFlux2TransformerTensorParallel(Flux2TransformerTesterConfig, TensorPar """Tensor Parallel inference tests for Flux2 Transformer (CUDA/XPU multi-accelerator).""" -def make_neuron_tp_spec(): - """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`). - - Returns `(model_class, init_dict, cpu_inputs)`. Defined here so all Flux2-specific test data lives in this file - while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest of - the Flux2 tests. - """ - config = Flux2TransformerTesterConfig() - return Flux2Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") - - class TestFlux2TransformerTensorParallelTPU(Flux2TransformerTesterConfig, TensorParallelTPUTesterMixin): """Tensor Parallel inference test for Flux2 Transformer on TPU.""" @@ -190,26 +176,8 @@ def get_init_dict(self): return {**super().get_init_dict(), "num_attention_heads": self.tp_world_size} -@is_tensor_parallel -@require_torch_neuron -class TestFlux2TransformerTensorParallelNeuron: - """Tensor Parallel inference test for Flux2 Transformer on AWS Neuron. - - Neuron TP runs through `torchrun` with the `"neuron"` distributed backend, so it cannot use the - `torch.multiprocessing`/NCCL spawn path of `TensorParallelTesterMixin`. This launches the generic worker - with the Flux2 model spec (`make_neuron_tp_spec`); the worker asserts the sharded output matches a single-device - reference, and the test checks its exit code. - """ - - def test_tensor_parallel_neuron_inference(self): - worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py") - spec = "tests.models.transformers.test_models_transformer_flux2:make_neuron_tp_spec" - cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec] - result = subprocess.run(cmd, capture_output=True, text=True) - assert result.returncode == 0, ( - f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n" - f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" - ) +class TestFlux2TransformerTensorParallelNeuron(Flux2TransformerTesterConfig, TensorParallelNeuronTesterMixin): + """Tensor Parallel inference test for Flux2 Transformer on AWS Neuron.""" class TestFlux2TransformerLoRA(Flux2TransformerTesterConfig, LoraTesterMixin): diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 6210aecff87f..0c4d3abe60c8 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -12,11 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - -import os -import subprocess -import sys - import pytest import torch @@ -24,7 +19,7 @@ from diffusers.models.transformers.transformer_qwenimage import compute_text_seq_len_from_mask 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, torch_device from ..testing_utils import ( AttentionBackendTesterMixin, AttentionTesterMixin, @@ -36,6 +31,7 @@ LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + TensorParallelNeuronTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, TorchAoTesterMixin, @@ -320,37 +316,8 @@ def test_tensor_parallel_tpu_inference(self): super().test_tensor_parallel_tpu_inference(atol=2e-2, rtol=2e-2) -def make_neuron_tp_spec(): - """Model spec consumed by the generic Neuron TP worker (``_neuron_tp_worker.py``). - - Returns ``(model_class, init_dict, cpu_inputs)``. Defined here so all QwenImage-specific test data lives in this - file while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the rest - of the QwenImage tests. - """ - config = QwenImageTransformerTesterConfig() - return QwenImageTransformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") - - -@is_tensor_parallel -@require_torch_neuron -class TestQwenImageTransformerTensorParallelNeuron: - """Tensor Parallel inference test for QwenImage Transformer on AWS Neuron. - - Neuron TP runs through ``torchrun`` with the ``"neuron"`` distributed backend, so it cannot use the - ``torch.multiprocessing``/NCCL spawn path of ``TensorParallelTesterMixin``. This launches the generic worker - with the QwenImage model spec (``make_neuron_tp_spec``); the worker asserts the sharded output matches a - single-device reference, and the test checks its exit code. - """ - - def test_tensor_parallel_neuron_inference(self): - worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py") - spec = "tests.models.transformers.test_models_transformer_qwenimage:make_neuron_tp_spec" - cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec] - result = subprocess.run(cmd, capture_output=True, text=True) - assert result.returncode == 0, ( - f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n" - f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}" - ) +class TestQwenImageTransformerTensorParallelNeuron(QwenImageTransformerTesterConfig, TensorParallelNeuronTesterMixin): + """Tensor Parallel inference test for QwenImage Transformer on AWS Neuron.""" class TestQwenImageTransformerLoRA(QwenImageTransformerTesterConfig, LoraTesterMixin): diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index f6d0a0004d0c..e178b954bd1d 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -29,6 +29,8 @@ ModelTesterMixin, SingleFileTesterMixin, TaylorSeerCacheTesterMixin, + TensorParallelNeuronTesterMixin, + TensorParallelTesterMixin, TrainingTesterMixin, ) @@ -380,3 +382,13 @@ def pretrained_model_kwargs(self): @property def torch_dtype(self): return torch.bfloat16 + + +class TestQwenImage21TransformerTensorParallel(QwenImage21TransformerTesterConfig, TensorParallelTesterMixin): + """Tensor Parallel inference tests for QwenImage 2.1 Transformer (CUDA/XPU multi-accelerator).""" + + +class TestQwenImage21TransformerTensorParallelNeuron( + QwenImage21TransformerTesterConfig, TensorParallelNeuronTesterMixin +): + """Tensor Parallel inference test for QwenImage 2.1 Transformer on AWS Neuron."""