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

* draft:add neuron as a legit backend

* feat: neuron-specific changes in the pipeline

* tests: eager tests

* draft: start with tp for flux2

* fix: style

* fix:apr_02 beta

* feat:add wan

* fix:pixart

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

* test: pixart compile mode on neuron

* cleanup & fix style

* merge: another change

* review: cleanup+suggestions

* fix: CIs style

* tests: add test units for tp

* fix: in case of text-encoder(s) on CPU

* review:cleanup+add test

* fix: style

* doc: remove it for now

* Add from_single_file support for SkyReelsV2 and ChronoEdit transformers (#13946)

SkyReels-V2 and ChronoEdit are both built on Wan, and their transformers have
the same keys as WanTransformer3DModel, so they reuse
convert_wan_transformer_to_diffusers (like WanVACE / WanAnimate). This lets the
community GGUF builds load directly.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

* multi-GPU VAE Fix for Cosmos 3 (#13924)

fix(cosmos3): pin VAE latent norm buffers to encode output device

Under sharded placement (device_map="balanced"), vae.encode() runs on the
VAE's own device while the mean/inv_std buffers were pinned to x.device,
causing a cross-device RuntimeError. Compute raw_mu first, then pin the
normalization buffers to its device so all tensors share one device.

Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* docs: fix repeated word typo in set_timesteps docstring (#13876)

* docs: fix repeated word typo in set_timesteps docstring

Removed the duplicate word "schedule" from the docstring for the sigmas argument in EulerDiscreteScheduler.set_timesteps.

* Update scheduling_euler_discrete.py

* Apply style fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

* clean some stuff to simplify code.

* clean more to remove permutation related shenanigans.

* revert: put torch.chunk back

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

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

* Address Sayak's TP review nits (docstrings + stale comments)

- _modeling_parallel: fix stale comment claiming attention processors read
  tp_degree at runtime (no longer true after the PackedColwise/Rowwise
  refactor; processors are TP-agnostic via head_dim).
- tensor_parallel: single backticks instead of double in docstrings, add
  input/output examples to _blocks_to_block_sizes and _resolve_tp_plan, and
  drop the stale "applies the Flux2 fused-weight permutations" line from
  apply_tensor_parallel (permuters were removed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Fix stale permuter reference in Flux2ParallelSelfAttnProcessor comment

The proportional QKV/MLP slice is now guaranteed by PackedColwiseParallel
block-by-block sharding, not the removed fused-weight permuter.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Derive TP backend inside apply_tensor_parallel from the mesh device type

Drops the `backend` parameter: apply_tensor_parallel already holds the TP
mesh, so it derives neuron-vs-default from `config._mesh.device_type` itself.
The caller no longer computes it. Documents why torch_device can't be used
(Neuron reports as "cpu" via torch's accelerator API). Removes the now-unused
is_torch_neuronx_available import from modeling_utils.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Simplify TP backend derivation to the mesh device type alone

Drop the redundant is_torch_neuronx_available() guard: a DeviceMesh("neuron", ...)
only exists when running on Neuron, so tp_mesh.device_type is the single source
of truth.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Remove incorrect "Neuron reports as cpu" rationale from TP backend docstring

Verified on the current native torch_neuronx drop that Neuron does surface as
the torch accelerator (torch._C._get_accelerator().type == "neuron",
get_device() == "neuron"); the old "reports as cpu" claim was stale. The
mesh-device-type derivation is unchanged — it's the device being sharded onto.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop stale "XLA" labels from Neuron device/runtime comments

Neuron is a native torch backend now (not XLA-routed): remove the "(XLA)"
parenthetical from the randn-on-CPU note in torch_utils and the "XLA runtime"
mention from the Neuron TP test docstring. Behavior unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* make style: reflow tensor_parallel docstrings

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Revert "revert: put torch.chunk back"

This reverts commit 7ea75f771aee1c16d5a5668343a92ea2717fb84a.

* fix: sharding for neuron + validate flux 1

* fix: change for check_repository_consistency

* feat: support qwen image as well

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

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

* review: keep bria and nucleusmoe out

* review: keep bria and nucleusmoe out

* review: keep bria and nucleusmoe out

* review: address agent review

* review: apply suggestions

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

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

* review: change example tp size to 4

* review: add flux tp test

* review: address comments

* review: improve the doc

* review:apply qwen rotary_emb per device

* review:remove redundant device def after PR #14383

* review: doc title suggestion

* review: restore copy comment

* review: remove cp divisibility explanation from doc on parallelism guide

* fix style

---------

Co-authored-by: Hz_Zhang <47402297+HaozheZhang6@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Atharva Joshi <atjoshi@nvidia.com>
Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Ramkumar R <ramkumarashvanth09@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>