[core] Tensor parallelism for Qwen-Image-2.1 by JingyaHuang · Pull Request #14865 · huggingface/diffusers

JingyaHuang · 2026-09-24T16:21:23Z

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

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 · 2026-09-24T16:27:02Z

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 · 2026-09-30T06:00:34Z

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


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

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 <noreply@anthropic.com>

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!

… device

DN6

Thank you!

…ffusers into add-qwen-image-21-tp-plan

sayakpaul · 2026-10-14T14:45:03Z

-        self.freqs = [freq.to(device) for freq in self.freqs]
+        freqs = self._get_device_freqs(device)


I think @DN6 meant we kept self.freqs = [freq.to(device) for freq in self.freqs] as is instead of this change.

Yeah, actually the new Neuron SDK drop(drop_oct) now supports complex tensors, so the default complex RoPE now runs on Neuron as is. I reverted the Neuron-specific workaround in this branch, and will open a new PR for reverting our workaround for support tp for Qwen-image.

#15042 to revert neuron RoPE

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.

sayakpaul · 2026-10-15T10:32:21Z

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


What is this change now?

with the latest nightly built of torch_tpu, torch.get_device_module("tpu").device_count() now raises
(RuntimeError: device_count is internal to TorchTPU and can only be called by the PyTorch framework (torch and torch_tpu). Use torch.accelerator.device_count instead.).

this is how it is working with beta releases... but torch.accelerator is equivalent and work for cuda/neuron at least.

It's available from PyTorch 2.6 as well because that is the minimum PyTorch we require.

github-actions Bot added documentation Improvements or additions to documentation lora models tests pipelines hooks size/L PR with diff > 200 LOC labels Sep 24, 2026

JingyaHuang force-pushed the add-qwen-image-21-tp-plan branch from cc669f0 to 12e2b9c Compare September 25, 2026 14:16

JingyaHuang marked this pull request as ready for review September 25, 2026 14:17

JingyaHuang requested a review from sayakpaul September 29, 2026 15:20

sayakpaul reviewed Sep 30, 2026

View reviewed changes

JingyaHuang and others added 4 commits October 2, 2026 16:07

feat: qwen image 21 tp plan

95c69bf

test: vallidated on neuron

481f25d

fix: lru_cache_unless_export leak

46b415a

fix: use complex RoPE on TPU for Qwen Image 2.1

f679fe5

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 <noreply@anthropic.com>

JingyaHuang force-pushed the add-qwen-image-21-tp-plan branch from fb74251 to f679fe5 Compare October 2, 2026 16:15

github-actions Bot added size/M PR with diff < 200 LOC and removed documentation Improvements or additions to documentation lora pipelines hooks size/L PR with diff > 200 LOC labels Oct 2, 2026

JingyaHuang and others added 3 commits October 7, 2026 14:31

refactor: inline the Neuron device check for Qwen Image 2.1 RoPE

0da8358

Merge branch 'main' into add-qwen-image-21-tp-plan

d641013

review: cache Qwen Image 2.1 RoPE device freqs with lru_cache

58a87bd

JingyaHuang requested a review from sayakpaul October 7, 2026 23:49

sayakpaul reviewed Oct 8, 2026

View reviewed changes

Comment thread src/diffusers/models/transformers/transformer_qwenimage21.py Outdated

sayakpaul reviewed Oct 8, 2026

View reviewed changes

review: drop redundant torch.device

1b294bb

Merge branch 'main' into add-qwen-image-21-tp-plan

281c3c5

sayakpaul requested a review from DN6 October 8, 2026 10:02

sayakpaul approved these changes Oct 8, 2026

View reviewed changes

Merge branch 'main' into add-qwen-image-21-tp-plan

f9c39bf

DN6 reviewed Oct 12, 2026

View reviewed changes

Comment thread tests/models/transformers/test_models_transformer_qwenimage21.py Outdated

Comment thread src/diffusers/models/transformers/transformer_qwenimage21.py Outdated

JingyaHuang and others added 5 commits October 12, 2026 15:39

Merge branch 'main' into add-qwen-image-21-tp-plan

6961b73

fix: use torch.accelerator.device_count() when resolving the parallel…

b301314

… device

refactor: drop the Neuron branch in Qwen Image 2.1 block ids

17c2aa8

restore: neuron repeat_interleave lowering error fix

615cd92

review: add TensorParallelNeuronTesterMixin for Neuron TP tests

aab59a2

github-actions Bot added size/L PR with diff > 200 LOC and removed size/M PR with diff < 200 LOC labels Oct 12, 2026

JingyaHuang requested review from DN6 and sayakpaul October 12, 2026 23:39

DN6 approved these changes Oct 14, 2026

View reviewed changes

Merge branch 'main' into add-qwen-image-21-tp-plan

9ac753c

sayakpaul approved these changes Oct 14, 2026

View reviewed changes

JingyaHuang and others added 3 commits October 14, 2026 14:30

Merge branch 'main' into add-qwen-image-21-tp-plan

593ada2

review: revert Qwen Image 2.1 RoPE freqs cache to a per-instance dict

6210f4c

Merge branch 'add-qwen-image-21-tp-plan' of github.com:JingyaHuang/di…

906b279

…ffusers into add-qwen-image-21-tp-plan

sayakpaul reviewed Oct 14, 2026

View reviewed changes

JingyaHuang and others added 2 commits October 15, 2026 10:40

Merge branch 'main' into add-qwen-image-21-tp-plan

46dc894

JingyaHuang requested a review from sayakpaul October 15, 2026 10:30

sayakpaul reviewed Oct 15, 2026

View reviewed changes

JingyaHuang mentioned this pull request Oct 15, 2026

refactor: drop the Neuron RoPE path in Qwen Image #15042

Merged

10 tasks

sayakpaul approved these changes Oct 15, 2026

View reviewed changes

sayakpaul merged commit ee84919 into huggingface:main Oct 15, 2026
56 of 58 checks passed

JingyaHuang deleted the add-qwen-image-21-tp-plan branch October 15, 2026 11:39

		self.freqs = [freq.to(device) for freq in self.freqs]
		freqs = self._get_device_freqs(device)

Conversation

JingyaHuang commented Sep 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Validation

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Sep 24, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

sayakpaul Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

DN6 left a comment

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

JingyaHuang Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JingyaHuang Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

JingyaHuang Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JingyaHuang commented Sep 24, 2026 •

edited

Loading

JingyaHuang Oct 15, 2026 •

edited

Loading

JingyaHuang Oct 15, 2026 •

edited

Loading