Jingya HUANG 1d5d056ecf [TPU] TorchTPU backend integration - eager / torch.compile / tp (#14039) ...

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

2026-10-09 16:16:23 +05:30

65 KiB

Raw Permalink Blame History

View Raw