[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

@@ -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
 

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:

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):