* feat: add torchtpu
* feat:draft TorchTPU support
* fix: wan overflow issue + compile mode error on sdxl
* doc: enhance with TorchTPU doc
* doc: enhance with TorchTPU doc
* style: remove unused imports flagged by ruff
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* docs: remove Debug Eager and Fused Eager sections from tpu.md
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* feat: add TP for TPU
* style: fix import sorting in TPU test scripts
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* style: ruff format TPU test scripts
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* fix: style
* fix: test for native 4 devices
* fix: propagate TPU device fixes to Flux/Flux2/Wan-family copies; register tpu.md; drop redundant execution_device check
- Propagate the text_encoder.device-based fix (introduced for TPU CPU-offload
support) from FluxPipeline/Flux2KleinPipeline/WanPipeline into their
`# Copied from` copies (flux/*, flux2_klein_inpaint, visualcloze, anyflow,
chronoedit, lucy_edit, skyreels_v2/*). SDXL-family copies of
StableDiffusionXLPipeline.encode_prompt are intentionally left untouched;
they'll be handled in a follow-up PR that fixes device placement for every
pipeline component (not just text encoders).
- Register docs/source/en/optimization/tpu.md in _toctree.yml (was breaking
the docs build: "not present in the table of contents").
- Remove the redundant "prefer non-CPU, non-meta component" loop from
DiffusionPipeline._execution_device: PR #14383 already fixed this in
DiffusionPipeline.device, which _execution_device falls back to. Verified
on TPU hardware that _execution_device still resolves correctly for a
split-placement pipeline after the removal.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
* tests: cleanup
* fix: fix compile mode
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* test: remove flux2 e2e test
* doc: apply suggestions
* doc: apply suggestions
* review: remove monkey patch
* review: revert neuron-specific changes in the tests
* review: apply suggestions
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* doc: add tpu to tp doc
* review: remove unnecessary for tpu
* test: assert every _tp_plan parameter is sharded after a TP load
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* removal: delete redundant tp shard helpers
* removal: drop TPU workarounds no longer needed after recent main changes
* doc: keep the FLUX.2-dev text encoder off a single TPU chip in the TP example
* doc: encode the prompt on CPU in the TPU TP example so it runs end to end
* test: tighten TPU TP tolerance and simplify the TPU TP worker
* doc: use enable_model_cpu_offload() in the TPU eager example
* doc: shard both FLUX.2-dev text encoder and transformer in the TPU TP example
* doc: run the text encoders on TPU in the compiled example
* test: drop the redundant TPU sync in the TP worker
* test: trim comments in the TPU TP tests
* test: run the TPU TP test with mp.spawn like the CUDA one
* test: pin the TPU TP test to 4 chips
* Update docs/source/en/optimization/tpu.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* doc: use DistributedConfig and dtype, and drop the torch_tpu import in the TPU examples
* test: skip the TPU TP test for models without a _tp_plan
* doc: correct the warning of recompilation
* fix: clone the SD3 timestep so torch.compile doesn't recompile every step, thus we can use it for the tpu.md compile mode snippet
* doc: update advice for compile mode when one chip doesn't fit
* test: pass the TPU TP tolerances as test arguments
* review: apply suggestions to the doc
* review: view compilation patch for tpu-only
---------
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: JingyaHuang <huang_jingya@outlok.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* 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>
