From 95c69bfb778ac60b38d7e6e40c13523a47b4b547 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 24 Sep 2026 16:15:08 +0000 Subject: [PATCH 01/13] feat: qwen image 21 tp plan --- .../transformers/transformer_qwenimage21.py | 85 ++++++++++++++++--- 1 file changed, 74 insertions(+), 11 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index b6eabfd584bb..4d850a0dc2d3 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import math from typing import Any @@ -23,7 +24,7 @@ from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import logging from ...utils.peft_utils import apply_lora_scale -from ...utils.torch_utils import maybe_allow_in_graph +from ...utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -133,6 +134,42 @@ def apply_rotary_emb_qwen( return x_out.type_as(x) +# Copied from diffusers.models.transformers.transformer_qwenimage.apply_rotary_emb_qwen_neuron +def apply_rotary_emb_qwen_neuron(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """ + Apply rotary embeddings to `x` using real-valued cos/sin, for backends without a complex dtype. + + Numerically equivalent to `apply_rotary_emb_qwen(..., use_real=False)`, which multiplies `x` by a complex + exponential. Neuron has no complex tensor support, so the rotation angles are carried as reals and cos/sin are + taken here instead. + + Args: + x (`torch.Tensor`): Query or key tensor to rotate, shape `[B, S, H, D]`. + freqs (`torch.Tensor`): Rotation angles, shape `[S, D // 2]`. + + Returns: + `torch.Tensor`: `x` with rotary embeddings applied. + """ + # Adjacent feature pairs (2k, 2k+1) share angle k, so each angle is repeated twice along the last dim; unsqueeze + # the head axis so the freqs broadcast over heads (this is what keeps it tensor-parallel-agnostic). + cos = torch.cos(freqs).repeat_interleave(2, dim=-1).unsqueeze(1) # [S, 1, D] + sin = torch.sin(freqs).repeat_interleave(2, dim=-1).unsqueeze(1) # [S, 1, D] + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) # [B, S, H, D] + return (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + + +# RoPE application is backend-dependent: the default path multiplies by a complex exponential, which Neuron and TPU +# cannot represent. On those backends `QwenImage21Rope` hands out rotation angles instead of complex freqs, and +# `apply_rotary_emb_qwen_neuron` takes cos/sin on device. Callers select by `device.type` and fall back to the +# default for any backend not listed here. +_ROPE_ANGLE_DEVICES = ("neuron", "tpu") +ROPE_PER_DEVICE = { + "cuda": functools.partial(apply_rotary_emb_qwen, use_real=False), + **dict.fromkeys(_ROPE_ANGLE_DEVICES, apply_rotary_emb_qwen_neuron), +} + + class QwenImage21TemporalTimesteps(nn.Module): r"""Sinusoidal timestep embedding. `cos` occupies the first half of the channels and `sin` the second.""" @@ -337,16 +374,20 @@ 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) if rotary_emb is not None: - query = apply_rotary_emb_qwen(query, rotary_emb, use_real=False) - key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) + apply_rope = ROPE_PER_DEVICE.get(query.device.type, ROPE_PER_DEVICE["cuda"]) + query = apply_rope(query, rotary_emb) + key = apply_rope(key, rotary_emb) if layer_cache is not None: if kv_cache_mode == "extract" and cache_write_slice is not None: @@ -674,10 +715,19 @@ def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torc freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) return torch.polar(torch.ones_like(freqs), freqs) + @lru_cache_unless_export(maxsize=None) + def _get_device_freqs(self, device: torch.device) -> list[torch.Tensor]: + """Return the per-axis freqs on `device`: complex exponentials, or rotation angles where complex is missing.""" + if device.type in _ROPE_ANGLE_DEVICES: + # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless because + # only cos/sin of the angle are used. + return [torch.angle(freq).to(device) for freq in self.freqs] + return [freq.to(device) for freq in self.freqs] + def forward( self, img_shapes: list[tuple[int, int, int]], image_pad_mask: torch.Tensor, device: torch.device ) -> torch.Tensor: - self.freqs = [freq.to(device) for freq in self.freqs] + freqs = self._get_device_freqs(torch.device(device)) frame_index, height_index, width_index = [], [], [] image_height_index, image_width_index = [], [] @@ -707,7 +757,7 @@ def forward( height_index[image_pad_mask] = torch.tensor(image_height_index, dtype=torch.long, device=device) width_index[image_pad_mask] = torch.tensor(image_width_index, dtype=torch.long, device=device) - return torch.cat([self.freqs[0][frame_index], self.freqs[1][height_index], self.freqs[2][width_index]], dim=-1) + return torch.cat([freqs[0][frame_index], freqs[1][height_index], freqs[2][width_index]], dim=-1) class QwenImage21Transformer2DModel( @@ -758,6 +808,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 +895,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), + # Built from the Python block lengths rather than with a tensor-repeats `repeat_interleave`, whose + # data-dependent output size some compiled backends (e.g. Neuron) cannot lower. + 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 From 481f25dc86b89dbf7c2554866856c5751cc51dc2 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 24 Sep 2026 16:18:09 +0000 Subject: [PATCH 02/13] test: vallidated on neuron --- .../test_models_transformer_qwenimage21.py | 43 ++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index f6d0a0004d0c..f22a8a1addef 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -13,6 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import subprocess +import sys + import pytest import torch from torch.nn.attention.flex_attention import create_mask @@ -21,7 +25,7 @@ from diffusers.models.transformers.transformer_qwenimage21 import build_qwenimage21_block_causal_mask from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -29,6 +33,7 @@ ModelTesterMixin, SingleFileTesterMixin, TaylorSeerCacheTesterMixin, + TensorParallelTesterMixin, TrainingTesterMixin, ) @@ -380,3 +385,39 @@ 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).""" + + +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)``. Reuses the shared tester config so the spec never drifts from the + rest of the QwenImage 2.1 tests. + """ + config = QwenImage21TransformerTesterConfig() + return QwenImage21Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") + + +@is_tensor_parallel +@require_torch_neuron +class TestQwenImage21TransformerTensorParallelNeuron: + """Tensor Parallel inference test for QwenImage 2.1 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 2.1 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_qwenimage21: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}" + ) From 46b415ab96bacd7f8b9d97cd989b9639eb462b43 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Fri, 25 Sep 2026 12:43:33 +0000 Subject: [PATCH 03/13] fix: lru_cache_unless_export leak --- .../transformers/transformer_qwenimage21.py | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 4d850a0dc2d3..98884c116ac7 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -24,7 +24,7 @@ from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import logging from ...utils.peft_utils import apply_lora_scale -from ...utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph +from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -161,9 +161,9 @@ def apply_rotary_emb_qwen_neuron(x: torch.Tensor, freqs: torch.Tensor) -> torch. # RoPE application is backend-dependent: the default path multiplies by a complex exponential, which Neuron and TPU # cannot represent. On those backends `QwenImage21Rope` hands out rotation angles instead of complex freqs, and -# `apply_rotary_emb_qwen_neuron` takes cos/sin on device. Callers select by `device.type` and fall back to the -# default for any backend not listed here. -_ROPE_ANGLE_DEVICES = ("neuron", "tpu") +# `apply_rotary_emb_qwen_neuron` takes cos/sin on device. Callers select by `device.type` (PyTorch/XLA reports TPU +# tensors as `"xla"`) and fall back to the default for any backend not listed here. +_ROPE_ANGLE_DEVICES = ("neuron", "xla") ROPE_PER_DEVICE = { "cuda": functools.partial(apply_rotary_emb_qwen, use_real=False), **dict.fromkeys(_ROPE_ANGLE_DEVICES, apply_rotary_emb_qwen_neuron), @@ -710,19 +710,24 @@ def __init__(self, theta: int, axes_dim: list[int]): torch.cat([self.rope_params(pos_index, dim, theta), self.rope_params(neg_index, dim, theta)], dim=0) for dim in axes_dim ] + # Per-device copies of `freqs`, kept on the instance so they are freed with the model. A class-level + # `lru_cache` would key on `self` and keep every instance's device freqs alive for the life of the process. + self._device_freqs: dict[torch.device, list[torch.Tensor]] = {} def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torch.Tensor: freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) return torch.polar(torch.ones_like(freqs), freqs) - @lru_cache_unless_export(maxsize=None) def _get_device_freqs(self, device: torch.device) -> list[torch.Tensor]: """Return the per-axis freqs on `device`: complex exponentials, or rotation angles where complex is missing.""" - if device.type in _ROPE_ANGLE_DEVICES: - # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless because - # only cos/sin of the angle are used. - return [torch.angle(freq).to(device) for freq in self.freqs] - return [freq.to(device) for freq in self.freqs] + if device not in self._device_freqs: + if device.type in _ROPE_ANGLE_DEVICES: + # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless + # because only cos/sin of the angle are used. + self._device_freqs[device] = [torch.angle(freq).to(device) for freq in self.freqs] + else: + self._device_freqs[device] = [freq.to(device) for freq in self.freqs] + return self._device_freqs[device] def forward( self, img_shapes: list[tuple[int, int, int]], image_pad_mask: torch.Tensor, device: torch.device From f679fe59b590b7f4426e68e9777940867861eb2b Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Fri, 25 Sep 2026 14:09:43 +0000 Subject: [PATCH 04/13] fix: use complex RoPE on TPU for Qwen Image 2.1 TorchTPU reports TPU tensors as "tpu" and supports complex dtypes, so only Neuron needs the angle-based RoPE path. Co-Authored-By: Claude Opus 5.5 --- .../models/transformers/transformer_qwenimage21.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 98884c116ac7..4d303c9b61e4 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -159,11 +159,11 @@ def apply_rotary_emb_qwen_neuron(x: torch.Tensor, freqs: torch.Tensor) -> torch. return (x.float() * cos + x_rotated.float() * sin).to(x.dtype) -# RoPE application is backend-dependent: the default path multiplies by a complex exponential, which Neuron and TPU -# cannot represent. On those backends `QwenImage21Rope` hands out rotation angles instead of complex freqs, and -# `apply_rotary_emb_qwen_neuron` takes cos/sin on device. Callers select by `device.type` (PyTorch/XLA reports TPU -# tensors as `"xla"`) and fall back to the default for any backend not listed here. -_ROPE_ANGLE_DEVICES = ("neuron", "xla") +# RoPE application is backend-dependent: the default path multiplies by a complex exponential, which Neuron cannot +# represent. On those backends `QwenImage21Rope` hands out rotation angles instead of complex freqs, and +# `apply_rotary_emb_qwen_neuron` takes cos/sin on device. Callers select by `device.type`. Every other backend (CUDA, +# CPU, TPU, …) uses the default complex path; only backends without complex dtypes are listed here. +_ROPE_ANGLE_DEVICES = ("neuron",) ROPE_PER_DEVICE = { "cuda": functools.partial(apply_rotary_emb_qwen, use_real=False), **dict.fromkeys(_ROPE_ANGLE_DEVICES, apply_rotary_emb_qwen_neuron), From 0da8358083758814fa5cb8e582fa06977c90b743 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Wed, 7 Oct 2026 14:31:42 +0000 Subject: [PATCH 05/13] refactor: inline the Neuron device check for Qwen Image 2.1 RoPE --- .../models/transformers/transformer_qwenimage21.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 4d303c9b61e4..ae2e62e1b067 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -160,13 +160,10 @@ def apply_rotary_emb_qwen_neuron(x: torch.Tensor, freqs: torch.Tensor) -> torch. # RoPE application is backend-dependent: the default path multiplies by a complex exponential, which Neuron cannot -# represent. On those backends `QwenImage21Rope` hands out rotation angles instead of complex freqs, and -# `apply_rotary_emb_qwen_neuron` takes cos/sin on device. Callers select by `device.type`. Every other backend (CUDA, -# CPU, TPU, …) uses the default complex path; only backends without complex dtypes are listed here. -_ROPE_ANGLE_DEVICES = ("neuron",) +# represent. Callers select by `device.type` and fall back to the default for any backend not listed here. ROPE_PER_DEVICE = { "cuda": functools.partial(apply_rotary_emb_qwen, use_real=False), - **dict.fromkeys(_ROPE_ANGLE_DEVICES, apply_rotary_emb_qwen_neuron), + "neuron": apply_rotary_emb_qwen_neuron, } @@ -721,7 +718,7 @@ def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torc def _get_device_freqs(self, device: torch.device) -> list[torch.Tensor]: """Return the per-axis freqs on `device`: complex exponentials, or rotation angles where complex is missing.""" if device not in self._device_freqs: - if device.type in _ROPE_ANGLE_DEVICES: + if device.type == "neuron": # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless # because only cos/sin of the angle are used. self._device_freqs[device] = [torch.angle(freq).to(device) for freq in self.freqs] From 58a87bd0b28ebe14a3e2978e21d0f8be1d73a139 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Wed, 7 Oct 2026 23:49:15 +0000 Subject: [PATCH 06/13] review: cache Qwen Image 2.1 RoPE device freqs with lru_cache --- .../transformers/transformer_qwenimage21.py | 35 ++++++++++--------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index ae2e62e1b067..57e8e32db35d 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -24,7 +24,7 @@ from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import logging from ...utils.peft_utils import apply_lora_scale -from ...utils.torch_utils import maybe_allow_in_graph +from ...utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -707,24 +707,19 @@ def __init__(self, theta: int, axes_dim: list[int]): torch.cat([self.rope_params(pos_index, dim, theta), self.rope_params(neg_index, dim, theta)], dim=0) for dim in axes_dim ] - # Per-device copies of `freqs`, kept on the instance so they are freed with the model. A class-level - # `lru_cache` would key on `self` and keep every instance's device freqs alive for the life of the process. - self._device_freqs: dict[torch.device, list[torch.Tensor]] = {} def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torch.Tensor: freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) return torch.polar(torch.ones_like(freqs), freqs) + @lru_cache_unless_export(maxsize=128) def _get_device_freqs(self, device: torch.device) -> list[torch.Tensor]: """Return the per-axis freqs on `device`: complex exponentials, or rotation angles where complex is missing.""" - if device not in self._device_freqs: - if device.type == "neuron": - # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless - # because only cos/sin of the angle are used. - self._device_freqs[device] = [torch.angle(freq).to(device) for freq in self.freqs] - else: - self._device_freqs[device] = [freq.to(device) for freq in self.freqs] - return self._device_freqs[device] + if device.type == "neuron": + # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless + # because only cos/sin of the angle are used. + return [torch.angle(freq).to(device) for freq in self.freqs] + return [freq.to(device) for freq in self.freqs] def forward( self, img_shapes: list[tuple[int, int, int]], image_pad_mask: torch.Tensor, device: torch.device @@ -897,11 +892,17 @@ def build_token_metadata( ) image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long) - # Built from the Python block lengths rather than with a tensor-repeats `repeat_interleave`, whose - # data-dependent output size some compiled backends (e.g. Neuron) cannot lower. - block_ids = torch.tensor( - [block for block, length in enumerate(block_lengths) for _ in range(length)], device=image_pad_mask.device - ) + if image_pad_mask.device.type == "neuron": + # 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, + ) + else: + block_ids = torch.repeat_interleave( + torch.arange(len(block_lengths), device=image_pad_mask.device), + torch.tensor(block_lengths, device=image_pad_mask.device), + ) image_ids[image_positions] = block_ids target_token_mask = torch.zeros_like(image_pad_mask) From 1b294bbe58af0ffd676c21f3a2fe820180d05977 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 8 Oct 2026 10:00:43 +0000 Subject: [PATCH 07/13] review: drop redundant torch.device --- src/diffusers/models/transformers/transformer_qwenimage21.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 57e8e32db35d..249198f041ae 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -724,7 +724,7 @@ def _get_device_freqs(self, device: torch.device) -> list[torch.Tensor]: def forward( self, img_shapes: list[tuple[int, int, int]], image_pad_mask: torch.Tensor, device: torch.device ) -> torch.Tensor: - freqs = self._get_device_freqs(torch.device(device)) + freqs = self._get_device_freqs(device) frame_index, height_index, width_index = [], [], [] image_height_index, image_width_index = [], [] From b301314cfd5cbb104bcd24f1bda5765a36fc386f Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Mon, 12 Oct 2026 15:08:16 +0000 Subject: [PATCH 08/13] fix: use torch.accelerator.device_count() when resolving the parallel device --- src/diffusers/models/modeling_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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: From 17c2aa8e2bd94cc78328a805c8e5a3cce36e5df7 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Mon, 12 Oct 2026 16:22:32 +0000 Subject: [PATCH 09/13] refactor: drop the Neuron branch in Qwen Image 2.1 block ids --- .../transformers/transformer_qwenimage21.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 249198f041ae..02037215190a 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -892,17 +892,12 @@ def build_token_metadata( ) image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long) - if image_pad_mask.device.type == "neuron": - # 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, - ) - else: - block_ids = torch.repeat_interleave( - torch.arange(len(block_lengths), device=image_pad_mask.device), - torch.tensor(block_lengths, device=image_pad_mask.device), - ) + # `output_size` fixes the output length up front, so the op does not depend on tensor values. + block_ids = torch.repeat_interleave( + torch.arange(len(block_lengths), device=image_pad_mask.device), + torch.tensor(block_lengths, device=image_pad_mask.device), + output_size=sum(block_lengths), + ) image_ids[image_positions] = block_ids target_token_mask = torch.zeros_like(image_pad_mask) From 615cd92a5119d1f05afb62e6dc43aadc163eda07 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Mon, 12 Oct 2026 21:47:20 +0000 Subject: [PATCH 10/13] restore: neuron repeat_interleave lowering error fix --- .../models/transformers/transformer_qwenimage21.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 02037215190a..5489faff0a5f 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -892,11 +892,10 @@ def build_token_metadata( ) image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long) - # `output_size` fixes the output length up front, so the op does not depend on tensor values. - block_ids = torch.repeat_interleave( - torch.arange(len(block_lengths), device=image_pad_mask.device), - torch.tensor(block_lengths, device=image_pad_mask.device), - output_size=sum(block_lengths), + # 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 From aab59a2374d24da2fc150c5f6c6b76f92e7515d6 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Mon, 12 Oct 2026 23:33:30 +0000 Subject: [PATCH 11/13] review: add TensorParallelNeuronTesterMixin for Neuron TP tests --- tests/models/testing_utils/__init__.py | 2 + tests/models/testing_utils/parallelism.py | 55 +++++++++ .../models/transformers/_neuron_tp_worker.py | 111 ------------------ .../test_models_transformer_flux.py | 39 +----- .../test_models_transformer_flux2.py | 40 +------ .../test_models_transformer_qwenimage.py | 41 +------ .../test_models_transformer_qwenimage21.py | 41 +------ 7 files changed, 75 insertions(+), 254 deletions(-) delete mode 100644 tests/models/transformers/_neuron_tp_worker.py 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 31d47f11535a..1c079ac4bbd6 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, @@ -44,6 +40,7 @@ MemoryTesterMixin, ModelTesterMixin, SingleFileTesterMixin, + TensorParallelNeuronTesterMixin, TensorParallelTesterMixin, TensorParallelTPUTesterMixin, TorchAoCompileTesterMixin, @@ -170,17 +167,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.""" @@ -189,26 +175,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 f22a8a1addef..e178b954bd1d 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.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 from torch.nn.attention.flex_attention import create_mask @@ -25,7 +21,7 @@ from diffusers.models.transformers.transformer_qwenimage21 import build_qwenimage21_block_causal_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 ( AttentionTesterMixin, BaseModelTesterConfig, @@ -33,6 +29,7 @@ ModelTesterMixin, SingleFileTesterMixin, TaylorSeerCacheTesterMixin, + TensorParallelNeuronTesterMixin, TensorParallelTesterMixin, TrainingTesterMixin, ) @@ -391,33 +388,7 @@ class TestQwenImage21TransformerTensorParallel(QwenImage21TransformerTesterConfi """Tensor Parallel inference tests for QwenImage 2.1 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)``. Reuses the shared tester config so the spec never drifts from the - rest of the QwenImage 2.1 tests. - """ - config = QwenImage21TransformerTesterConfig() - return QwenImage21Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") - - -@is_tensor_parallel -@require_torch_neuron -class TestQwenImage21TransformerTensorParallelNeuron: - """Tensor Parallel inference test for QwenImage 2.1 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 2.1 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_qwenimage21: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 TestQwenImage21TransformerTensorParallelNeuron( + QwenImage21TransformerTesterConfig, TensorParallelNeuronTesterMixin +): + """Tensor Parallel inference test for QwenImage 2.1 Transformer on AWS Neuron.""" From 6210f4c4da596c05928722401529411db492e6e9 Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Wed, 14 Oct 2026 13:36:56 +0000 Subject: [PATCH 12/13] review: revert Qwen Image 2.1 RoPE freqs cache to a per-instance dict --- .../transformers/transformer_qwenimage21.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index 5489faff0a5f..a6dffe316a17 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -24,7 +24,7 @@ from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import logging from ...utils.peft_utils import apply_lora_scale -from ...utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph +from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin @@ -707,19 +707,24 @@ def __init__(self, theta: int, axes_dim: list[int]): torch.cat([self.rope_params(pos_index, dim, theta), self.rope_params(neg_index, dim, theta)], dim=0) for dim in axes_dim ] + # Per-device copies of `freqs`, kept on the instance so they are freed with the model. A class-level + # `lru_cache` would key on `self` and keep every instance's device freqs alive for the life of the process. + self._device_freqs: dict[torch.device, list[torch.Tensor]] = {} def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torch.Tensor: freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) return torch.polar(torch.ones_like(freqs), freqs) - @lru_cache_unless_export(maxsize=128) def _get_device_freqs(self, device: torch.device) -> list[torch.Tensor]: """Return the per-axis freqs on `device`: complex exponentials, or rotation angles where complex is missing.""" - if device.type == "neuron": - # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless - # because only cos/sin of the angle are used. - return [torch.angle(freq).to(device) for freq in self.freqs] - return [freq.to(device) for freq in self.freqs] + if device not in self._device_freqs: + if device.type == "neuron": + # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless + # because only cos/sin of the angle are used. + self._device_freqs[device] = [torch.angle(freq).to(device) for freq in self.freqs] + else: + self._device_freqs[device] = [freq.to(device) for freq in self.freqs] + return self._device_freqs[device] def forward( self, img_shapes: list[tuple[int, int, int]], image_pad_mask: torch.Tensor, device: torch.device From bed111599b0e89dd5c4a031b3666e873ac483e6f Mon Sep 17 00:00:00 2001 From: JingyaHuang Date: Thu, 15 Oct 2026 10:07:50 +0000 Subject: [PATCH 13/13] review: drop the Neuron RoPE path in Qwen Image 2.1 Neuron now runs complex tensors natively, so Qwen Image 2.1 uses the default complex RoPE on every device. This removes `apply_rotary_emb_qwen_neuron`, `ROPE_PER_DEVICE` and the per-device freqs cache, and keeps `self.freqs = [freq.to(device) ...]` as on main. --- .../transformers/transformer_qwenimage21.py | 57 ++----------------- 1 file changed, 4 insertions(+), 53 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index a6dffe316a17..1a9ef250c519 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/src/diffusers/models/transformers/transformer_qwenimage21.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import functools import math from typing import Any @@ -134,39 +133,6 @@ def apply_rotary_emb_qwen( return x_out.type_as(x) -# Copied from diffusers.models.transformers.transformer_qwenimage.apply_rotary_emb_qwen_neuron -def apply_rotary_emb_qwen_neuron(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: - """ - Apply rotary embeddings to `x` using real-valued cos/sin, for backends without a complex dtype. - - Numerically equivalent to `apply_rotary_emb_qwen(..., use_real=False)`, which multiplies `x` by a complex - exponential. Neuron has no complex tensor support, so the rotation angles are carried as reals and cos/sin are - taken here instead. - - Args: - x (`torch.Tensor`): Query or key tensor to rotate, shape `[B, S, H, D]`. - freqs (`torch.Tensor`): Rotation angles, shape `[S, D // 2]`. - - Returns: - `torch.Tensor`: `x` with rotary embeddings applied. - """ - # Adjacent feature pairs (2k, 2k+1) share angle k, so each angle is repeated twice along the last dim; unsqueeze - # the head axis so the freqs broadcast over heads (this is what keeps it tensor-parallel-agnostic). - cos = torch.cos(freqs).repeat_interleave(2, dim=-1).unsqueeze(1) # [S, 1, D] - sin = torch.sin(freqs).repeat_interleave(2, dim=-1).unsqueeze(1) # [S, 1, D] - x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] - x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) # [B, S, H, D] - return (x.float() * cos + x_rotated.float() * sin).to(x.dtype) - - -# RoPE application is backend-dependent: the default path multiplies by a complex exponential, which Neuron cannot -# represent. Callers select by `device.type` and fall back to the default for any backend not listed here. -ROPE_PER_DEVICE = { - "cuda": functools.partial(apply_rotary_emb_qwen, use_real=False), - "neuron": apply_rotary_emb_qwen_neuron, -} - - class QwenImage21TemporalTimesteps(nn.Module): r"""Sinusoidal timestep embedding. `cos` occupies the first half of the channels and `sin` the second.""" @@ -382,9 +348,8 @@ def _qwenimage21_prepare_qkv( key = attn.norm_k(key).to(value.dtype) if rotary_emb is not None: - apply_rope = ROPE_PER_DEVICE.get(query.device.type, ROPE_PER_DEVICE["cuda"]) - query = apply_rope(query, rotary_emb) - key = apply_rope(key, rotary_emb) + query = apply_rotary_emb_qwen(query, rotary_emb, use_real=False) + key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) if layer_cache is not None: if kv_cache_mode == "extract" and cache_write_slice is not None: @@ -707,29 +672,15 @@ def __init__(self, theta: int, axes_dim: list[int]): torch.cat([self.rope_params(pos_index, dim, theta), self.rope_params(neg_index, dim, theta)], dim=0) for dim in axes_dim ] - # Per-device copies of `freqs`, kept on the instance so they are freed with the model. A class-level - # `lru_cache` would key on `self` and keep every instance's device freqs alive for the life of the process. - self._device_freqs: dict[torch.device, list[torch.Tensor]] = {} def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torch.Tensor: freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) return torch.polar(torch.ones_like(freqs), freqs) - def _get_device_freqs(self, device: torch.device) -> list[torch.Tensor]: - """Return the per-axis freqs on `device`: complex exponentials, or rotation angles where complex is missing.""" - if device not in self._device_freqs: - if device.type == "neuron": - # `torch.angle` runs on CPU while the freqs are still complex; wrapping into (-pi, pi] is harmless - # because only cos/sin of the angle are used. - self._device_freqs[device] = [torch.angle(freq).to(device) for freq in self.freqs] - else: - self._device_freqs[device] = [freq.to(device) for freq in self.freqs] - return self._device_freqs[device] - def forward( self, img_shapes: list[tuple[int, int, int]], image_pad_mask: torch.Tensor, device: torch.device ) -> torch.Tensor: - freqs = self._get_device_freqs(device) + self.freqs = [freq.to(device) for freq in self.freqs] frame_index, height_index, width_index = [], [], [] image_height_index, image_width_index = [], [] @@ -759,7 +710,7 @@ def forward( height_index[image_pad_mask] = torch.tensor(image_height_index, dtype=torch.long, device=device) width_index[image_pad_mask] = torch.tensor(image_width_index, dtype=torch.long, device=device) - return torch.cat([freqs[0][frame_index], freqs[1][height_index], freqs[2][width_index]], dim=-1) + return torch.cat([self.freqs[0][frame_index], self.freqs[1][height_index], self.freqs[2][width_index]], dim=-1) class QwenImage21Transformer2DModel(