* 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 7ea75f771a.
* 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>
* initial
* fix code quality
* dependency test fix
* apply review suggestions
* add extras and version check
* fix gating and prefer local cache for modular
* simplify transformers registration and make the version warning silenceable
* remove pipeline level test
* only model level testing
* layers to quantizze, change em dashes
* env flag for transformers
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Remove JAX/Flax support
Deletes all Flax models, schedulers, pipelines, and loading utilities,
the from_flax loading path, flax/jax packaging extras, Flax CI filters,
Flax training examples, and Flax content in translated docs.
Hub compatibility is preserved: Flax class names in old model_index.json
files are still stripped to their PyTorch equivalents, and *.msgpack
files in repos are still ignored during download.
* Drop Flax class-name remapping from model_index.json loading
Flax-saved checkpoint metadata is no longer recognized: class names like
FlaxStableDiffusionPipeline in model_index.json now fail to resolve
instead of being remapped to their PyTorch equivalents. *.msgpack files
in hub repos are still ignored during download.
* Restore Flax training examples
The example scripts are standalone and still work with diffusers<=0.39.x.
A warning note is added to each Flax README section pointing out the
version requirement now that the library itself no longer ships Flax.
* update
* Repoint scheduler default-config tests off the Flax fixture
* Repoint example training tests off the Flax fixture
* port final set of model tests and others
* fix extracter.
* fix cuda tests for models.
* Revert "fix extracter."
This reverts commit 178c4cb99b.
* Revert "port final set of model tests and others"
This reverts commit a92c70c081.
* fix more
* address reviewer feedback
---------
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* switch to transformers main again./
* more
* up
* up
* fix group offloading.
* attributes
* up
* up
* tie embedding issue.
* fix t5 stuff for more.
* matrix configuration to see differences between 4.57.3 and main failures.
* change qwen expected slice because of how init is handled in v5.
* same stuff.
* up
* up
* Revert "up"
This reverts commit 515dd06db5.
* Revert "up"
This reverts commit 5274ffdd7f.
* up
* up
* fix with peft_format.
* just keep main for easier debugging.
* remove torchvision.
* empty
* up
* up with skyreelsv2 fixes.
* fix skyreels type annotation.
* up
* up
* fix variant loading issues.
* more fixes.
* fix dduf
* fix
* fix
* fix
* more fixes
* fixes
* up
* up
* fix dduf test
* up
* more
* update
* hopefully ,final?
* one last breath
* always install from main
* up
* audioldm tests
* up
* fix PRX tests.
* up
* kandinsky fixes
* qwen fixes.
* prx
* hidream
* drop python 3.8
* remove list, tuple, dict from typing
* fold Unions into |
* up
* fix a bunch and please me.
* up
* up
* up
* up
* up
* up
* enforce 3.10.0.
* up
* up
* up
* up
* up
* up
* up
* up
* Update setup.py
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
* up.
* python 3.10.
* ifx
* up
* up
* up
* up
* final
* up
* fix typing utils.
* up
* up
* up
* up
* up
* up
* fix
* up
* up
* up
* up
* up
* up
* handle modern types.
* up
* up
* fix ip adapter type checking.
* up
* up
* up
* up
* up
* up
* up
* revert docstring changes.
* keep deleted files deleted.
* keep deleted files deleted.
---------
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
* fix: group offloading to support standalone computational layers in block-level offloading
* test: for models with standalone and deeply nested layers in block-level offloading
* feat: support for block-level offloading in group offloading config
* fix: group offload block modules to AutoencoderKL and AutoencoderKLWan
* fix: update group offloading tests to use AutoencoderKL and adjust input dimensions
* refactor: streamline block offloading logic
* Apply style fixes
* update tests
* update
* fix for failing tests
* clean up
* revert to use skip_keys
* clean up
---------
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
