diffusers

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

[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>