[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) device. Here as suggested we use Flux2 Klein as the starting point. But the TP support here is generic, easy to extend to other backend(cuda, tpu and more) and is exposed through the existing public API used for CP: model.enable_parallelism(config=TensorParallelConfig(...)).

Key changes:

• A model-agnostic apply_tensor_parallel that shards from a flat _tp_plan.

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]

With TP, context_parallel_config can be None, we set up _parallel_config._cp_world_size for it.

[sayakpaul]

Thanks for working on this.

There is lot of intrusive and model-specific changes which I think is a bit of an anti-pattern. I think it's also probably because of some of the fusion stuff that's happening inside Flux2.

More specifically, the intrusive pieces exist for one reason: Flux2 fuses projections into single Linears (SwiGLU gate+linear, and to_qkv_mlp_proj packing Q/K/V and MLP).

Contiguous column sharding is blind to that internal layout, so:

• you must reorder rows so each rank gets paired slices -> the permuters, and
• the local tensor width no longer factors as heads × head_dim or splits cleanly into qkv/mlp -> the runtime local_* recomputation.

I opened JingyaHuang#1 to simplify some of the stuff. LMK.

Furthermore, would the changes related to fusing be the same for Flux1, for example? I think gf the layers were unfused, parallelize_module + DTensor would handle head-splitting automatically and none of this would be needed.

[sayakpaul]

Can this not be derived from torch_device?

[sayakpaul]

Seems like it should be present in _modeling_parallel.py?

[sayakpaul]

Where is this needed?

[sayakpaul]

Is this not needed on CUDA?

[sayakpaul]

@JingyaHuang did my PR break any neuron-specific stuff?