[core] Tensor parallelism for Qwen-Image-2.1

[JingyaHuang]

What does this PR do?

Adds tensor-parallel support to QwenImage21Transformer2DModel, using the sharded checkpoint loading
from #14544.

Changes in `transformer_qwenimage21.p

• _tp_plan: attention to_q/to_k/to_v and SwiGLU proj/gate_layer are colwise; to_out.0
  and img_mlp.out are rowwise. The shections, norm_out and proj_out stay replicated.
• Head split by head_dim insteal works when each rank holds only partof the heads (same fix as Qwen-Image).
• RoPE on Neuron: Neuron has no coPE with real cos/sin(apply_rotary_emb_qwen_neuron, copied from Qwen-Image). All other devices keep the complex path. The
  per-device freqs are now cached inste` on every forward.
• block_ids on Neuron: Neuron can't lower a tensor-repeats repeat_interleave, so on Neuron the
  block ids are built from the Python bre unchanged. This was needed for 2.1to run on Neuron at all, TP or not.

Tests: TensorParallelTesterMixin for CUDA/XPU, and a Neuron TP=2 test that compares the sharded output
against a single-device reference.

import torch
from torch import distributed as dist
from diffusers import QwenImage21Pipeline, QwenImage21Transformer2DModel, TensorParallelConfig

dist.init_process_group(backend="nccl")
rank, world_size = dist.get_rank(), d
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)

transformer = QwenImage21Transformer2
    "Qwen/Qwen-Image-2.1",
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    parallel_config=TensorParallelCon
)
pipe = QwenImage21Pipeline.from_pretrtransformer=transformer,torch_dtype=torch.bfloat16)
pipe.text_encoder.to(device)
pipe.vae.to(device)

image = pipe(
    "A capybara wearing a wizard hat,
    generator=torch.Generator().manual_seed(0),  # same seed on every rank
).images[0]
if rank == 0:
    image.save("qwen21_tp.png")
dist.destroy_process_group()

Run with torchrun --nproc-per-node 8 qwen21_tp.py.

Validation

Mode	Neuron	TPU	CUDA
Eager	✅	✅	✅
Compile	✅		✅

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

Left some comments on the Qwen specific stuff.

[sayakpaul]

I don't think we need this kind of dict munging. Let's just do: "neuron": apply_rotary_emb_qwen_neuron.

[sayakpaul]

Let's keep it explicitly conditioned on neuron then. @DN6 WDYT?

[sayakpaul]

Left one comment.

Ran the tests on HF Jobs using 2 A10g:
https://huggingface.co/jobs/sayakpaul/6ac71e14df2184ac91ac74c6

Also, ran the snippet and it ran fine (using 8 GPUs):

(HF Jobs)

[sayakpaul]

Thanks! Just ran some tests and they're good. The changes in the DiT code don't lead to any changes in the outputs!