[Neuron] Add tensor parallel support for Neuron backend by JingyaHuang · Pull Request #13718 · huggingface/diffusers

JingyaHuang · 2026-05-11T14:17:44Z

What does this PR do?

Adds tensor-parallel (TP) inference for diffusers models on AWS Neuron (Trainium/Inferentia).
The implementation is:

model-agnostic, it shards from a flat _tp_plan
the TP support is generic, easy to extend to other backends (CUDA, TPU, and more). It is exposed through the same public API used for CP: model.enable_parallelism(config=TensorParallelConfig(...)).

Now validated on 3 pipelines on Neuron (trn2, TP=8), in both eager and torch.compile mode: FLUX.1-dev, FLUX.2 (Klein), and Qwen-Image.

Key changes:

A model-agnostic apply_tensor_parallel that shards from a flat _tp_plan (Neuron pre-shard path works around the NRT consecutive-reduce_scatter bug; the default parallelize_module path is used on other backends).
_tp_plan added to the FLUX.1, FLUX.2 and Qwen-Image transformers.
Qwen-Image: the attention processor reshape is made TP-agnostic (reshape by a fixed head_dim), and its RoPE is ported from complex torch.polar/view_as_complex to real cos/sin. The RoPE change is numerically identical and unconditional — required for XLA backends (Neuron/TPU) and cleaner under torch.compile. The same real-RoPE change is applied to NucleusMoE, which shared the code.

Example scripts

Runnable torchrun --nproc_per_node=8 scripts live under examples_tp/ for each pipeline (e.g. test_neuron_flux1_dev_tp.py, test_neuron_flux2_dev_tp.py, test_qwenimage_tp.py).

Quick test — Flux2 TP on Neuron (For future release)

run with torchrun --nproc_per_node=8 flux2_tp8_neuron.py

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import DeviceMesh
import torch_neuronx  # noqa: F401 — registers torch.neuron

from diffusers import Flux2KleinPipeline, TensorParallelConfig

MODEL = "black-forest-labs/FLUX.2-klein-9B"
PROMPT = "a golden retriever surfing a wave, photorealistic"

dist.init_process_group(backend="neuron")
device = torch.neuron.current_device()
rank = dist.get_rank()
tp_size = dist.get_world_size()
tp_mesh = DeviceMesh("neuron", list(range(tp_size)))

pipe = Flux2KleinPipeline.from_pretrained(MODEL, torch_dtype=torch.bfloat16)

# Text encoder + VAE: replicated on every rank (no TP).
pipe.text_encoder = pipe.text_encoder.to(device)
pipe.vae = pipe.vae.to(device)

# Transformer: shard across all ranks while still on CPU, then move to device.
pipe.transformer.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh))
pipe.transformer = pipe.transformer.to(device)
torch.neuron.synchronize()

image = pipe(
    prompt=PROMPT, height=1024, width=1024,
    num_inference_steps=4, guidance_scale=1.0,
  ).images[0]

if rank == 0:
    image.save("flux2_tp8.png")
    print("Saved flux2_tp8.png")

dist.destroy_process_group()

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.

… into add-neuron-backend

sayakpaul

Thanks a lot for the iteration, @JingyaHuang! Left some further comments.

I am also running the tests and trying out examples. Will keep this PR updated with findings from that.

Edit: Here's the report of my findings https://gist.github.com/sayakpaul/954b5d64aad648aca091f36c35f36397

sayakpaul · 2026-07-28T03:44:38Z

+
+## Tensor parallelism
+
+[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) shards the weight matrices of a model across devices. Each device holds a column-wise (`"colwise"`) or row-wise (`"rowwise"`) slice of each layer, computes a partial result, and an `AllReduce`/`AllGather` at the layer boundary reconstructs the full output. Unlike context parallelism, it reduces the per-device *weight* memory, which is useful for models that do not fit on a single device.


Can we also provide a short section after the TP section, showing which parallelism should be preferred in what circumstances?

Also, should there be guidance for model authors on how they should write the _tp_plan for their model? That seems non-trivial to me and some useful hints could be nice to have there.

For guidance on which parallelism config to use, I think each section already explains briefly the trade-offs , eg. CP -> long sequences, TP -> large weights. I just added a section to summarize i bit, If we want more precise/actionable guide, we need to run experiments, since the right choice also depends on the specific model.

Will add more explanation on contributing tp plan tho!

I think all of that makes sense and should be documented (for example, the right choices depending on the model size, input problem space, etc.).

sayakpaul · 2026-07-28T04:02:32Z

-    to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
-    reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
-    tensors contain rotary embeddings and are returned as real tensors.
+def apply_rotary_emb_qwen(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:


Why is this change required?

It's not a tp required change, cuda runs the complex path fine. But for neuron, rope_params did torch.polar, so pos_freqs/neg_freqs are complex64, the Neuron compiler has no lowering for complex tensors so far. It's a workaround.

I think this is a breaking change. So, we should perhaps have a pattern like:

ROPE_PER_DEVICE = { "cuda": apply_rotary_emb_qwen, "neuron": apply_rotary_emb_qwen_neuron, }

And then fetch from ROPE_PER_DEVICE in the caller site. @DN6 WDYT?

sayakpaul · 2026-07-28T04:41:48Z

@@ -63,7 +63,12 @@
 from ..utils.distributed_utils import is_torch_dist_rank_zero


I think the following aren't implemented at the moment (which is fine; just flagging).

Sharded loading. from_pretrained should stream shards straight to each rank's DTensor rather than materializing the full checkpoint then slicing — otherwise TP saves you nothing at load time. And check save_pretrained / state_dict calls .full_tensor() or uses DCP. I think we should at least raise when save_pretrained() is called in case TP is enabled?

LoRA loading. for a colwise base layer, lora_A replicated + lora_B colwise; for rowwise, lora_A rowwise + lora_B replicated. If the plan doesn't cover PEFT layers, loading an adapter onto a TP model will either error or be wrong. I think we should detect if the model has peft layers injected and raise if TP is requested?

Quantization, offloading. We should probably also raise when these are requested?

@JingyaHuang this doesn't seem to have been resolved?

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

…into support-neuron-tp

DN6

Looking good 👍🏽 I think we can merge quite soon. My comments are minor.
@sayakpaul Perhpas we can handle sharded TP loading/saving + LoRA and Quant support follow ups so that the scope is manageable here.

DN6 · 2026-07-31T10:41:11Z

+        # When text encoders are offloaded to CPU while the denoising backbone
+        # (transformer, unet, vae) runs on an accelerator, self.device returns CPU
+        # (first component). Prefer any non-CPU, non-meta component so that
+        # latent tensors land on the accelerator. This covers CUDA, XPU, NPU, HPU,
+        # and any other backend, including TP-sharded models via DTensor.
+        for name, model in self.components.items():
+            if isinstance(model, torch.nn.Module) and model.device.type not in ("cpu", "meta"):
+                return model.device


This fix is okay with me, but is there a specific reason to include it in this PR? Does it affect TP?

sayakpaul · 2026-08-01T00:24:11Z

Perhpas we can handle sharded TP loading/saving + LoRA and Quant support follow ups so that the scope is manageable here.

200 percent.

tengomucho

just few comments

sayakpaul · 2026-08-04T14:34:32Z

+
+`tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices.
+
+### Writing a _tp_plan


@stevhliu can this first _ be escaped so that it renders properly? If so, how?

sayakpaul · 2026-08-04T14:35:40Z

+# in the model's _tp_plan
+"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),
+```
+


(nit): Perhaps provide a reference implementation example of Flux2 (by linking the path to the modeling file)?

sayakpaul

Thanks a lot for iterating. Just a few remaining nits.

JingyaHuang and others added 25 commits March 18, 2026 11:15

draft:add neuron as a legit backend

98f6c8c

Merge branch 'huggingface:main' into add-neuron-backend

c58b8b8

Merge branch 'huggingface:main' into add-neuron-backend

3367409

Merge branch 'main' into add-neuron-backend

0c51734

feat: neuron-specific changes in the pipeline

a76953c

tests: eager tests

2480388

draft: start with tp for flux2

1469c04

fix: style

929ab72

Merge branch 'huggingface:main' into add-neuron-backend

52cac76

Merge branch 'huggingface:main' into support-neuron-tp

30cb353

Merge branch 'add-neuron-backend' of github.com:JingyaHuang/diffusers…

28a5086

… into add-neuron-backend

Merge branch 'huggingface:main' into support-neuron-tp

7fab0c4

Merge branch 'huggingface:main' into add-neuron-backend

68689e5

Merge branch 'main' into add-neuron-backend

da79308

fix:apr_02 beta

3bb9c7c

Merge branch 'add-neuron-backend' of github.com:JingyaHuang/diffusers…

c4facab

… into add-neuron-backend

feat:add wan

dff1f32

Merge branch 'huggingface:main' into support-neuron-tp

1c930c4

Merge branch 'huggingface:main' into add-neuron-backend

1eb5ff9

fix:pixart

cbe8f28

fix: rewrite flux swiglu activation to avoid gather op in neuron IR

16b9606

test: pixart compile mode on neuron

7f13f68

Merge branch 'main' into neuron-torch-comppile

a46cb19

cleanup & fix style

a354b88

Merge branch 'neuron-torch-comppile' into support-neuron-tp

931bb85

github-actions Bot added size/L PR with diff > 200 LOC lora models tests utils labels May 11, 2026

JingyaHuang and others added 11 commits July 10, 2026 15:51

Merge branch 'main' into support-neuron-tp

6d84695

review: keep bria and nucleusmoe out

c6e0680

review: keep bria and nucleusmoe out

1646814

review: keep bria and nucleusmoe out

2218453

review: address agent review

3613423

Merge branch 'main' into support-neuron-tp

e8dcefc

review: apply suggestions

d55ac18

Merge branch 'main' into support-neuron-tp

48ea484

Merge branch 'main' into support-neuron-tp

fcc2bcb

Merge branch 'main' into support-neuron-tp

2364649

Merge branch 'main' into support-neuron-tp

ad5536d

sayakpaul reviewed Jul 28, 2026

View reviewed changes

Comment thread docs/source/en/training/distributed_inference.md Outdated

sayakpaul reviewed Jul 28, 2026

View reviewed changes

Comment thread docs/source/en/training/distributed_inference.md

sayakpaul reviewed Jul 28, 2026

View reviewed changes

JingyaHuang and others added 4 commits July 28, 2026 12:17

Update docs/source/en/training/distributed_inference.md

737945a

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

review: change example tp size to 4

686e379

Merge branch 'support-neuron-tp' of github.com:JingyaHuang/diffusers …

5c2549d

…into support-neuron-tp

review: add flux tp test

57365c5

DN6 reviewed Jul 31, 2026

View reviewed changes

tengomucho reviewed Aug 3, 2026

View reviewed changes

Comment thread src/diffusers/hooks/tensor_parallel.py

Comment thread src/diffusers/hooks/tensor_parallel.py

Comment thread src/diffusers/models/_modeling_parallel.py Outdated

JingyaHuang and others added 3 commits August 4, 2026 10:02

review: address comments

3c80014

Merge branch 'main' into support-neuron-tp

31112cb

Merge branch 'main' into support-neuron-tp

6f3db51

sayakpaul reviewed Aug 4, 2026

View reviewed changes

Comment thread docs/source/en/training/distributed_inference.md

sayakpaul reviewed Aug 4, 2026

View reviewed changes

JingyaHuang mentioned this pull request Aug 4, 2026

[core] DiffusionPipeline.device deduction for split-device pipelines #14383

Open


		## Tensor parallelism

		[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) shards the weight matrices of a model across devices. Each device holds a column-wise (`"colwise"`) or row-wise (`"rowwise"`) slice of each layer, computes a partial result, and an `AllReduce`/`AllGather` at the layer boundary reconstructs the full output. Unlike context parallelism, it reduces the per-device weight memory, which is useful for models that do not fit on a single device.

		@@ -63,7 +63,12 @@
		from ..utils.distributed_utils import is_torch_dist_rank_zero


		`tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices.

		### Writing a _tp_plan

Uh oh!

Conversation

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

Uh oh!

What does this PR do?

Key changes:

Example scripts

Quick test — Flux2 TP on Neuron (For future release)

Who can review?

Uh oh!

sayakpaul left a comment • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sayakpaul Jul 28, 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!

Choose a reason for hiding this comment

Uh oh!

DN6 left a comment • 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!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul commented Aug 1, 2026

Uh oh!

tengomucho left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

JingyaHuang commented May 11, 2026 •

edited

Loading

sayakpaul left a comment •

edited

Loading

sayakpaul Jul 28, 2026 •

edited

Loading

DN6 left a comment •

edited

Loading