[TPU] TorchTPU backend integration - eager / torch.compile / tp by JingyaHuang · Pull Request #14039 · huggingface/diffusers

JingyaHuang · 2026-06-22T14:50:43Z

What does this PR do?

Need the fix #14739 first.

Add TorchTPU support: eager mode + torch.compile
Add tests for TPU
Some doc / examples with TorchTPU
Following the addition of tensor parallelism in [core] Support tensor parallelism for model inference (CUDA, Neuron) #13718, this PR extends the support to TPU.

This is a preparation based on TorchTPU beta before the official release.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…ster 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 huggingface#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>

HuggingFaceDocBuilderDev · 2026-09-08T13:18:59Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

sayakpaul · 2026-09-08T20:33:05Z

+pipe.enable_tpu_compile()
+
+# Warmup — triggers static graph compilation.
+pipe.tpu_warmup(


tpu_warmup -- where is this function coming from?

sayakpaul · 2026-09-08T20:34:58Z

+"""TPU backend for tensor parallelism, dispatched from ``apply_tensor_parallel(backend="tpu")``.
+
+The structure mirrors the Neuron backend (``tensor_parallel_neuron.py``). The motivation differs: on Neuron the
+pre-shard path works around an NRT consecutive-reduce-scatter bug; here it prevents OOM. Without pre-sharding,
+``parallelize_module`` calls ``distribute_tensor`` internally, which loads the full weight matrix on every TPU chip
+before scattering. For large diffusion models this exhausts HBM. Pre-sharding each weight on CPU via
+``DTensor.from_local`` first means each chip only receives its local shard, then ``parallelize_module`` is called as a
+no-op for weights (they are already DTensors) but still registers the input/output hooks for the forward pass.


But this also seems like a generally useful thing for me to do on GPUs as well. Why keep this separate?

Yeah, this PR was opened in parallel with the TP initial support PR, it would be better to get #14544 landed first, so that the shards are tackled in a more generic way on TPU.

sayakpaul · 2026-09-08T20:38:40Z

    reference, and the test checks its exit code.
    """

+    WORLD_SIZE = 2


Can keep the Neuron related changes in a separate PR please?

sayakpaul · 2026-09-08T20:40:51Z

+
+@is_tensor_parallel
+@require_torch_tpu
+class TestFlux2TransformerTensorParallelTPU:


Can we not create a tester mixin class for this like

diffusers/tests/models/testing_utils/parallelism.py

Line 300 in c5469b7

class TensorParallelTesterMixin:

and use that instead?

It follows how we test other parallelism techniques.

stevhliu

nice, thanks for adding docs!

do we need to also mention TP here or no?

stevhliu · 2026-09-08T23:56:10Z

+> [!TIP]
+> For the best production throughput, prefer `torch.compile` via `pipe.enable_tpu_compile()`.
+
+## API reference


the API section should gets it own separate page under API > Main classes under Parallel inference as thats the closest related topic

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

…rs into add-torchtpu-support

JingyaHuang added 3 commits May 29, 2026 18:15

feat: add torchtpu

38007b4

feat:draft TorchTPU support

8ed15f9

fix: wan overflow issue + compile mode error on sdxl

e343c0a

github-actions Bot added documentation Improvements or additions to documentation models utils pipelines size/L PR with diff > 200 LOC labels Jun 22, 2026

JingyaHuang and others added 14 commits June 25, 2026 14:16

doc: enhance with TorchTPU doc

84b4049

doc: enhance with TorchTPU doc

bb3ec1e

Merge branch 'main' into add-torchtpu-support

339be41

Merge branch 'main' into add-torchtpu-support

92193a7

Merge branch 'main' into add-torchtpu-support

1a2369b

style: remove unused imports flagged by ruff

66254b4

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

docs: remove Debug Eager and Fused Eager sections from tpu.md

7c241be

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Merge branch 'main' into add-torchtpu-support

05d56d5

Merge branch 'main' into add-torchtpu-support

b7a8c0b

Merge branch 'huggingface:main' into add-torchtpu-support

b0c5595

Merge branch 'main' into add-torchtpu-support

7c3df6c

feat: add TP for TPU

7c1e379

style: fix import sorting in TPU test scripts

f944c48

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

style: ruff format TPU test scripts

82d20ab

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

JingyaHuang changed the title ~~[TPU] Initial TorchTPU backend integration (eager + torch.compile)~~ [TPU] TorchTPU backend integration - eager / torch.compile / tp Sep 7, 2026

github-actions Bot added tests hooks labels Sep 7, 2026

JingyaHuang mentioned this pull request Sep 7, 2026

[TPU] Add tensor parallel support for TPU backend #14075

Closed

JingyaHuang and others added 2 commits September 7, 2026 10:22

fix: style

f6c17d2

fix: test for native 4 devices

e8fe48c

JingyaHuang marked this pull request as ready for review September 7, 2026 16:23

github-actions Bot added the modular-pipelines label Sep 8, 2026

tests: cleanup

744fa95

JingyaHuang requested a review from sayakpaul September 8, 2026 16:42

sayakpaul reviewed Sep 8, 2026

View reviewed changes

fix: fix compile mode

bfc7d17

stevhliu reviewed Sep 9, 2026

View reviewed changes

JingyaHuang and others added 17 commits September 9, 2026 17:35

Update docs/source/en/optimization/tpu.md

10d4ad1

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

a999a3d

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

cb2dcf3

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Merge branch 'main' into add-torchtpu-support

9f11d1e

Update docs/source/en/optimization/tpu.md

c48d41c

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

aa70a45

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

6cb37d2

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

3e06f71

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

09d5ce2

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

7b46e08

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

e82e795

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

Update docs/source/en/optimization/tpu.md

d2bf529

Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>

test: remove flux2 e2e test

22d595a

doc: apply suggestions

042a88a

doc: apply suggestions

c6a4c80

Merge branch 'add-torchtpu-support' of github.com:JingyaHuang/diffuse…

c82fce7

…rs into add-torchtpu-support

review: remove monkey patch

cb314bd

github-actions Bot removed the models label Sep 10, 2026

review: revert neuron-specific changes in the tests

bf8934f

Uh oh!

Conversation

JingyaHuang commented Jun 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Sep 8, 2026

Uh oh!

Uh oh!

Uh oh!

sayakpaul Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

JingyaHuang Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

sayakpaul Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

stevhliu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

stevhliu Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JingyaHuang commented Jun 22, 2026 •

edited

Loading