[core] Tensor parallelism for Qwen-Image-2.1

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:

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,39 @@ 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."""
 
@@ -337,16 +371,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 +712,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=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]
+
     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(device)
 
         frame_index, height_index, width_index = [], [], []
         image_height_index, image_width_index = [], []
@@ -707,7 +754,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 +805,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 +892,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
 

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",

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: