* feat: qwen image 21 tp plan
* test: vallidated on neuron
* fix: lru_cache_unless_export leak
* fix: use complex RoPE on TPU for Qwen Image 2.1
TorchTPU reports TPU tensors as "tpu" and supports complex dtypes, so
only Neuron needs the angle-based RoPE path.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* refactor: inline the Neuron device check for Qwen Image 2.1 RoPE
* review: cache Qwen Image 2.1 RoPE device freqs with lru_cache
* review: drop redundant torch.device
* fix: use torch.accelerator.device_count() when resolving the parallel device
* refactor: drop the Neuron branch in Qwen Image 2.1 block ids
* restore: neuron repeat_interleave lowering error fix
* review: add TensorParallelNeuronTesterMixin for Neuron TP tests
* review: revert Qwen Image 2.1 RoPE freqs cache to a per-instance dict
* review: drop the Neuron RoPE path in Qwen Image 2.1
Neuron now runs complex tensors natively, so Qwen Image 2.1 uses the
default complex RoPE on every device. This removes
`apply_rotary_emb_qwen_neuron`, `ROPE_PER_DEVICE` and the per-device
freqs cache, and keeps `self.freqs = [freq.to(device) ...]` as on main.
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* 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>
* [core] Shard tensor-parallel checkpoints on load and save
Stream each rank's slice of a tensor-parallel checkpoint straight off disk
instead of materializing the full checkpoint on every rank and resharding it
afterwards, and gather the shards back on save.
- `from_pretrained(..., parallel_config=TensorParallelConfig(...))` resolves
the shard specs on the still-meta model, then slices each safetensors tensor
before the dtype cast, so host memory peaks at ~1/tp_degree of the checkpoint.
- `save_pretrained` all-gathers the DTensors into an ordinary checkpoint, or
writes a distributed checkpoint with `dcp=True` so no full tensor is ever
formed. The writing `tp_degree` is recorded, since a packed weight's stored
layout is interleaved by it.
- Factor the plan interpretation out of the Neuron pre-shard path into shared
`TPShardSpec` / `resolve_tp_shard_specs` / `_local_shard` / `_hooks_only_styles`
helpers, so both backends and both the load and save paths shard identically.
* Raise when tensor parallelism is combined with quantization, offloading or LoRA
Addresses the remaining two items of the review on #13718: tensor parallelism was rejected
alongside quantization and `device_map` only on the `from_pretrained` streaming path, while
`enable_parallelism` — which the quantization error message itself recommended — accepted a
quantized, offloaded or adapter-injected model and sharded it anyway.
- Add `_check_tp_model_state`, called from `apply_tensor_parallel`, the one chokepoint every TP
entry point funnels through. It rejects a model that is quantized, group-offloaded, placed by
accelerate (`device_map` or CPU offload), or has PEFT layers injected. Placed before the
device-type check so the reported reason is the useful one.
- Guard the reverse order too: `enable_group_offload`, the two pipeline CPU-offload methods, and
`load_lora_adapter` now refuse a tensor-parallel model.
- `save_pretrained` refuses a quantized tensor-parallel model. Previously the `dcp=True` branch
returned before the quantizer's serialization step, writing shards with no quantization
metadata and no error.
- The DCP load guard checked the `quantization_config` kwarg only, so a pre-quantized checkpoint
directory loaded silently; check the config's own entry too, and add the missing `_tp_plan`
check that otherwise surfaced as a raw `AttributeError`.
- Correct the `from_pretrained` message and the doc sentence that pointed at `enable_parallelism`
as a way to shard a quantized model.
The new tests are the first tensor-parallel tests that need neither an accelerator nor more than
one rank: every case asserts a raise before any collective, so they run single-process on gloo.
* review: remove tp save/ dcp related
* review: add a new helper .
* review: moved the TP checks into enable_parallelism
* review: move the warning under caution block
* review: add loading time numbers
* review: move tp size dividende check
* review: delete test
* review: consolidate tp checks
* review: consolidate the tp checks, all in tensor_parallel.py
* review: improve conditionals
* review: debug info
* review: check weights format before instantiating the model
* review: restore duplicated _find_mismatched_keys since irrelevant
* revert: delete duplicated helper
* review: update doc for loading benchmark of flux2-dev
* review: improve the raising check
* review: add an args section
* Test tensor-parallel loading from sharded checkpoints
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.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>
* add `_flash_3_varlen_hub` support for SP
* make signiture clear
* update
* add accuracy test
* fix dtype
* loosen threshold because of bfloat16
* fix the break change of https://github.com/huggingface/kernels-community/commit/1471f94
* fix stale error message
* fix merge
* revert change
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* add mask support for flash backend
* fix test case
* refactor test
* add protection
* fix comment
* update according to suggestion
* revert change
* fix according to claude review
* add test converage for QwenImage
* add SP support and fix non-contiguous mask for flash_varlen kernel
* revert change
* Update tests/models/testing_utils/parallelism.py
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Update tests/models/testing_utils/parallelism.py
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* drop `_padded_to_unpad`
* follow `if _parallel_config is None` pattern
* rename `attn_mask_2d`
* move check to the top
* make comment clear
* move non-contiguous-attention-mask as default dummy data
* revert and update
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* tests: add cp backend and attention backend tests.
* up
* up
* up
* fix ring for flash and flash_3
* generate.
* Apply suggestions from code review
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
* up
* up
---------
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
* fix mask in SP
* change the modification to qwen specific
* drop xfail since qwen-image mask is fixed
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* UT expands to batch inputs
* update according to suggestion
* update according to suggestion 2
* fix CI
* update according to suggestion 3
* clean line
