[core] Support tensor parallelism for model inference (CUDA, Neuron)

[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

importtorchimporttorch.distributedasdistfromtorch.distributed.device_meshimportDeviceMeshimporttorch_neuronx# noqa: F401 — registers torch.neuronfromdiffusersimportFlux2KleinPipeline, TensorParallelConfigMODEL="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]
ifrank==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.

[sayakpaul]

@zhtmike could you check this PR and provide feedback? Additionally, I think your Slack connection got expired. Could you let me know what is the best email to reach you?

[zhtmike]

Hi @sayakpaul, got your ping! I’ve sent it to your email.

My major concern is that this PR changes the QKV layout for five models, and there is no regression test to ensure the model output is the same before and after this change. It may break the results when the attention backend is set or when CP is enabled.

The other changes look fine to me, with minor suggestions.

[zhtmike]

I suggest adding a Mixin to guard the functionality and accuracy when both context parallel and tensor parallel are enabled, since they are orthogonal. Maybe in a future PR?

[sayakpaul]

Thanks @zhtmike!

My major concern is that this PR changes the QKV layout for five models, and there is no regression test to ensure the model output is the same before and after this change. It may break the results when the attention backend is set or when CP is enabled.

Do you have a skeleton for such a test? We can run it quickly and see if this is concerning?

[zhtmike]

Thanks @zhtmike!

My major concern is that this PR changes the QKV layout for five models, and there is no regression test to ensure the model output is the same before and after this change. It may break the results when the attention backend is set or when CP is enabled.

Do you have a skeleton for such a test? We can run it quickly and see if this is concerning?

Just my opinion: a more conservative approach is to keep the model’s default layout unchanged and perform the permutation in the TP block (in column/row parallel if necessary).

Otherwise, I suggest simply running some inference and checking the results by eye. It’s hard to compare the old and new layouts in a simple test — may need a shell script to switch branches for this kind of testing...

[zhtmike]

Thanks @zhtmike!

My major concern is that this PR changes the QKV layout for five models, and there is no regression test to ensure the model output is the same before and after this change. It may break the results when the attention backend is set or when CP is enabled.

Do you have a skeleton for such a test? We can run it quickly and see if this is concerning?

Just my opinion: a more conservative approach is to keep the model’s default layout unchanged and perform the permutation in the TP block (in column/row parallel if necessary).

Otherwise, I suggest simply running some inference and checking the results by eye. It’s hard to compare the old and new layouts in a simple test — may need a shell script to switch branches for this kind of testing...

An update, the layout seems not changed. (I overlooked attn.heads, attn.head_dim as same thing...) Yeah, so the PR looks no problem for me. Running some simple inference test might be enough.

[sayakpaul]

Used the following (Claude generated; but verified by me) script and I think it should address the concern around anything being broken around attention backends:

Details

<script>Unfold</script>

importargparseimportosimportsysp=argparse.ArgumentParser()
p.add_argument("--repo", default=".")
p.add_argument("--backends", nargs="*", default=[])
p.add_argument("--save", default=None)
p.add_argument("--ref", default=None)
args=p.parse_args()
sys.path[:0] = [args.repo, os.path.join(args.repo, "src")]
importtorchfromdiffusers.models.attention_dispatchimport (
_HUB_KERNELS_REGISTRY,
_AttentionBackendRegistry,
AttentionBackendName,
)
# flash-attn2 hub kernel v1 has no torch 2.13 builds; v3 works for the single-GPU pathforbin (AttentionBackendName.FLASH_HUB, AttentionBackendName.FLASH_VARLEN_HUB):
cfg=_HUB_KERNELS_REGISTRY[b]
cfg.version, cfg.wrapped_forward_attr, cfg.wrapped_backward_attr=3, None, Nonefromtests.models.transformers.test_models_transformer_fluximportFluxTransformerTesterConfigfromtests.models.transformers.test_models_transformer_flux2importFlux2TransformerTesterConfigfromtests.models.transformers.test_models_transformer_qwenimageimportQwenImageTransformerTesterConfigtorch.use_deterministic_algorithms(False)
calls= {}
forname, fninlist(_AttentionBackendRegistry._backends.items()):
defwrap(f, n):
definner(*a, **k):
calls[n] =calls.get(n, 0) +1returnf(*a, **k)
returninner_AttentionBackendRegistry._backends[name] =wrap(fn, name.value)
definputs_for(name, cfg):
n=32*32ifname=="flux":
inputs=cfg.get_dummy_inputs()
inputs["hidden_states"] =torch.randn(1, n, 4, generator=cfg.generator)
inputs["img_ids"] =torch.randn(n, 3, generator=cfg.generator)
elifname=="flux2":
inputs=cfg.get_dummy_inputs(height=32, width=32)
else:
inputs=cfg.get_dummy_inputs()
inputs["hidden_states"] =torch.randn(1, n, 16, generator=cfg.generator)
inputs["img_shapes"] = [(1, 32, 32)]
inputs["encoder_hidden_states_mask"] =Nonereturn {k: v.cpu() iftorch.is_tensor(v) elsevfork, vininputs.items()}
defcast(inputs, dtype):
return {
k: v.to("cuda", dtype) iftorch.is_tensor(v) andv.is_floating_point()
elsev.to("cuda") iftorch.is_tensor(v) elsevfork, vininputs.items()
}
CONFIGS= {
"flux": FluxTransformerTesterConfig,
"flux2": Flux2TransformerTesterConfig,
"qwenimage": QwenImageTransformerTesterConfig,
}
forname, ConfiginCONFIGS.items():
tc=Config()
torch.manual_seed(0)
model=tc.model_class(**tc.get_init_dict())
ifargs.ref:
model.load_state_dict(torch.load(f"{args.ref}/{name}_sd.pt", weights_only=True))
inputs=torch.load(f"{args.ref}/{name}_inputs.pt", weights_only=False)
else:
inputs=inputs_for(name, tc)
model=model.to("cuda").eval()
withtorch.no_grad():
out=model(**cast(inputs, torch.float32), return_dict=False)[0].cpu()
ifargs.ref:
d= (out-torch.load(f"{args.ref}/{name}_out.pt", weights_only=True)).abs()
print(f"{name}: vs ref (fp32 native) max={d.max():.3e} mean={d.mean():.3e}")
ifargs.save:
os.makedirs(args.save, exist_ok=True)
torch.save({k: v.cpu() fork, vinmodel.state_dict().items()}, f"{args.save}/{name}_sd.pt")
torch.save(inputs, f"{args.save}/{name}_inputs.pt")
torch.save(out, f"{args.save}/{name}_out.pt")
ifargs.backends:
model=model.to(torch.bfloat16)
binputs=cast(inputs, torch.bfloat16)
withtorch.no_grad():
native=model(**binputs, return_dict=False)[0]
restore, _=_AttentionBackendRegistry.get_active_backend()
forbackendinargs.backends:
model.set_attention_backend(backend)
calls.clear()
withtorch.no_grad():
out_b=model(**binputs, return_dict=False)[0]
model.reset_attention_backend()
_AttentionBackendRegistry.set_active_backend(restore)
d= (out_b-native).abs()
ok=torch.allclose(out_b, native, atol=1e-2, rtol=1e-2)
print(f" {name}/{backend}: max={d.max():.3e} mean={d.mean():.3e} ok={ok} dispatched={dict(calls)}")

Compare across different backends against SDPA:

python backend_check.py --backends _native_cudnn flash_hub _flash_3_hub

[sayakpaul]

Thanks a lot for working on this. Let's make sure to open a follow-up for the stuff we listed here so that the erroring is better for our users.