[Neuron] Add tensor parallel support for Neuron backend

[JingyaHuang]

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.

[JingyaHuang]

Made some modifications according to your comments @sayakpaul, and added flux1 support to test if your PR can easily adapt to suppot it!

[JingyaHuang]

in https://github.com/JingyaHuang/diffusers/blob/44eba6edd78181cb7e20a1585c5b86679b9e75ad/src/diffusers/models/attention_dispatch.py#L1120

if grad_enabled or (_parallel_config is not None and _parallel_config.context_parallel_config._world_size > 1):

_parallel_config.context_parallel_config would be None when TP enabled, so the property was added to pass the conditional statements

[JingyaHuang]

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py

Line 419 in d1b988b

out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data)

this change is for dynamic index-driven scatter_ in _unpack_latents_with_ids, neuron compiler can't handle it, haven't tested on TPU yet but likely an issue as well. for cuda, scatter_ is a native on-device op, so it's ok.
actually for neuron etc, it would be even better to reformulate unpack_latents_with_ids so it doesn't do a data-dependent scatter to avoid the device-host sync penalty...

[sayakpaul]

Made some modifications according to your comments @sayakpaul, and added flux1 support to test if your PR can easily adapt to suppot it!

Does it? 👀👀

[sayakpaul]

Let's keep Bria out?

[JingyaHuang]

To keep Bria and nucleusmoe out, I would need to break the copied from links, eg:

diffusers/src/diffusers/models/transformers/transformer_bria_fibo.py

Lines 67 to 68 in 284419b

	# Copied from diffusers.models.transformers.transformer_flux.FluxAttnProcessor with FluxAttnProcessor->BriaFiboAttnProcessor, FluxAttention->BriaFiboAttention
	class BriaFiboAttnProcessor:

[sayakpaul]

Interesting. But should self.device not return CPU in this because the entire pipeline isn't on an accelerator?

[JingyaHuang]

true, right now self.device is define by the 1st component in alphabetical order, which is not ideal. But how would you define self.device if the components of the pipeline are in different devices? shall we define it by a certain model in the pipeline?

[JingyaHuang]

perhaps:

@property
def device(self) -> torch.device:
    r"""
    Returns:
        `torch.device`: The torch device on which the pipeline is located. When components are
        split across devices (e.g. text encoders on CPU while the denoising backbone runs on an
        accelerator), the accelerator device is returned.
    """
    module_names, _ = self._get_signature_keys(self)
    modules = [getattr(self, n, None) for n in module_names]
    modules = [m for m in modules if isinstance(m, torch.nn.Module)]

    # Prefer a non-CPU, non-meta component so a split pipeline reports the accelerator
    # it computes on, rather than whichever component happens to sort first.
    for module in modules:
        if module.device.type not in ("cpu", "meta"):
            return module.device

    for module in modules:
        return module.device

    return torch.device("cpu")

?

[sergereview]

🤗 Serge says:

Solid, well-scoped TP implementation overall — the plan-driven sharding design, the packed colwise/rowwise handling for fused projections, and the real-valued RoPE rewrite (verified numerically identical to the complex formulation) all check out. However there are a few correctness issues that should be fixed before merge.

Correctness

• Unconditional per-step broadcast in pipeline_flux2_klein.py: it fires whenever torch.distributed is initialized, not just under TP. Any distributed non-TP use of this pipeline (e.g. data-parallel with different prompts/seeds per rank) will have every rank's latents silently overwritten by rank 0's. It is also self-described as a "safety measure only", which the repo rules explicitly disallow.
• Packed-colwise bias is sharded incorrectly in the generic path (hooks/tensor_parallel.py): _PackedColwiseImpl shards non-weight params with a plain contiguous Shard(0), but the weight rows are packed per-block. For any packed layer with a bias the bias slice won't line up with the weight shard. The Neuron path (_pre_shard_and_tp) packs the bias correctly — the generic path should match.
• Missed attention backend: _flash_attention_3_varlen_hub (attention_dispatch.py ~L3216) still dereferences _parallel_config.context_parallel_config.ring_degree and checks if _parallel_config is None: unguarded. With TP-only, _parallel_config is set but context_parallel_config is None, so this backend will crash with AttributeError. All sibling backends were updated in this PR; this one was skipped.
• No divisibility validation when computing per-rank chunks (chunk = bs // tp_size, rows = w.shape[0] // tp_size in both backends): if a block size isn't divisible by tp_degree, rows/columns are silently dropped and the model produces garbage. The head-divisibility check in enable_parallelism doesn't cover MLP dims. A concise ValueError would be much better than silent truncation.

API / scope concerns

• apply_rotary_emb_qwen / _apply_rotary_emb_nucleus signature change (drops use_real / use_real_unbind_dim, changes freqs semantics from complex to angles) is a breaking change to a shipped public-ish function; downstream code importing it will break. Worth an explicit sign-off from maintainers.
• The PipelineMixin.device / _execution_device changes alter long-standing semantics for every pipeline (any stray non-CPU component now defines the pipeline device). This is a broad behavioral change smuggled into a TP PR and deserves its own discussion/tests.
• The # Copied from ...apply_rotary_emb_qwen header on _apply_rotary_emb_nucleus was removed even though the two functions remain identical — keeping the link (and running make fix-copies) would prevent future drift.
• Duplicate CP+TP mutual-exclusion check: ParallelConfig.__post_init__ already raises, making the second check in enable_parallelism unreachable for normal construction.

Tests / description

• The PR description references runnable scripts under examples_tp/, but no such directory exists in the diff — the description should be updated.
• The generic TensorParallelTesterMixin and the Neuron torchrun worker are well structured; tolerances are justified.

serge v0.1.0 · model: claude-fable-5 · 18 LLM turns · 38 tool calls · 1236.4s · 1418671 in / 90227 out tokens

[sergereview]

Bug: for a packed-colwise layer with a bias, this shards the bias with a plain contiguous Shard(0), but the weight rows above are sharded per block. E.g. for blocks=[Q, K, V, gate, up], rank r's weight shard holds the r-th slice of each block, while its bias shard holds a contiguous 1/tp slice of the full fused bias — they don't correspond, so outputs are wrong whenever bias=True.

The Neuron path (_pre_shard_and_tp) packs the bias with the same per-block slicing as the weight (bias_parts loop); this generic path needs the same treatment. Currently only reachable with bias=True fused projections, but Flux2FeedForward/Flux2ParallelSelfAttention both expose a bias config knob, so this is a live footgun.

[sergereview]

This changes the documented semantics of pipeline.device for every pipeline, not just TP ones: any single non-CPU component (e.g. a VAE left on GPU after enable_model_cpu_offload finished a generation, or a partially-moved pipeline) now redefines the whole pipeline's reported device, where previously the first component won. That's a global behavior change to a core public property riding along in a TP PR, with no tests covering the new precedence. Consider scoping the TP fix to _execution_device only (where latents placement actually matters) and leaving device untouched, or splitting this into its own PR with tests.

[stevhliu]

very nice docs, thanks for adding!