[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 and perferrably merge the sharding improvement PR #14544 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>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… example

… end

sayakpaul

Nice, thanks! The changes are minimal and we should be able to merge soon.

(Ran the updated tests on 2 A10Gs and they worked: https://huggingface.co/jobs/sayakpaul/6ac30cecfbc85ba6823a4556)

sayakpaul · 2026-10-05T02:00:19Z

+| Mode | Constant | How to activate | Notes |
+|---|---|---|---|
+| Strict eager (default) | `EagerMode.DEFER_NEVER` | `import torch_tpu` | Operations dispatched one at a time, asynchronous |
+| Compile | — | `torch.compile(module, backend="tpu")` | AOT compilation with `TpuBackend` |


Should we present a faithful comparison between their latency numbers?

Not at this stage but maybe in the future, need to ask TPU team.

Sorry but why do we need to ask the TPU team about presenting compilation numbers?

sayakpaul · 2026-10-05T02:02:29Z

+image.save("output.png")
+```
+
+If the text encoder alone is too large for a single chip(eg. FLUX.2-dev's Mistral-3-Small is ~45GB),


I think we should provide a working example here. Or we could use transformers' support for TP?

Also, the same argument (a model to be placed on a TPU chip is too large) applies to the DiT. I think we could make an entirely separate section to talk about TP, etc.

Or, we could include a statement like "For details on how to shard the denoising module, refer to the "## Tensor parallelism" section."

Yeah, it would be better just show a simple example with everything ran on tpu w/o tp in this section, and have an example with TP on the TP dedicated section.

sayakpaul · 2026-10-05T02:09:31Z

+pipe.transformer = torch.compile(pipe.transformer, backend="tpu", fullgraph=True, dynamic=False)
+pipe.vae = torch.compile(pipe.vae, backend="tpu", fullgraph=True, dynamic=False)


Not a necessity but:

We can simplify this by calling the module-level compile method like transformer.compile().

Have we verified if compile_repeated_blocks() works on the transformer? It is related to regional compilation.

transformer.compile() with offload fails when setting fullgraph=True, because accelerate's CpuOffload.pre_forward is wrapped in torch.compiler.disable.
https://github.com/huggingface/accelerate/blob/7d8824f10f8be1ea405d611c98d51ebfbf4821cd/src/accelerate/hooks.py#L748-L757

And yes, compile_repeated_blocks() has been verified on the TPU. All components run on tpu:0, and the output is correct.

sayakpaul · 2026-10-05T02:19:56Z

-            sharded = {k: v for k, v in model.state_dict().items() if isinstance(v, DTensor)}
-            assert sharded, "No parameter was sharded into a DTensor by the streaming load."
-            name, param = next(iter(sharded.items()))
+            specs = resolve_tp_shard_specs(model, model_class._tp_plan, world_size)


This is so good! Thanks for the work, here!

… example

stevhliu · 2026-10-06T16:18:30Z

+> TorchTPU requires **static shapes** — pass `dynamic=False`. Every time `height`, `width`, or
+> `num_inference_steps` changes, the graph is recompiled from scratch. Keep these values constant
+> across all calls after warmup, or run another warmup pass before changing them.


the code example below compiles only the transformer's repeated blocks, so changing num_inference_steps shouldn't trigger a recompile?

stevhliu · 2026-10-06T16:24:53Z

+
+repo_id = "black-forest-labs/FLUX.2-dev"
+text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
+    repo_id, subfolder="text_encoder", dtype=torch.bfloat16, tp_plan="auto", device_mesh=mesh


i think passing tp_plan in from_pretrained is going to be removed soon, so it may be better to use DistributedConfig here instead

sayakpaul

Thanks! I left some further comments.

sayakpaul · 2026-10-07T05:55:07Z

+| Mode | Constant | How to activate | Notes |
+|---|---|---|---|
+| Strict eager (default) | `EagerMode.DEFER_NEVER` | `import torch_tpu` | Operations dispatched one at a time, asynchronous |
+| Compile | — | `torch.compile(module, backend="tpu")` | AOT compilation with `TpuBackend` |


Sorry but why do we need to ask the TPU team about presenting compilation numbers?

sayakpaul · 2026-10-07T06:23:10Z

+
+    def test_tensor_parallel_tpu_inference(self):
+        from torch_tpu._internal.distributed.launchers.singlehost_wrapper import prepare_tpu_environment
+        from torch_tpu._internal.utils import hardware


Should we not also skip the tests, if the underlying model class doesn't have any _tp_plan defined?

sayakpaul · 2026-10-07T06:24:27Z

+    tp_atol = 1e-3
+    tp_rtol = 1e-3


Should these be configurable from test_tensor_parallel_tpu_inference? Easy to override no?

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

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 models label Sep 19, 2026

Merge branch 'main' into add-torchtpu-support

a638798

github-actions Bot removed the modular-pipelines label Sep 21, 2026

JingyaHuang and others added 2 commits September 21, 2026 15:13

review: remove unnecessary for tpu

2bf3619

Merge branch 'main' into add-torchtpu-support

5d6ad22

JingyaHuang mentioned this pull request Oct 1, 2026

[core] Shard tensor-parallel checkpoints on load and save #14544

Merged

10 tasks done

JingyaHuang and others added 4 commits October 1, 2026 13:45

Merge branch 'main' into add-torchtpu-support

eb49a61

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

test: assert every _tp_plan parameter is sharded after a TP load

3152988

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

removal: delete redundant tp shard helpers

7257a79

removal: drop TPU workarounds no longer needed after recent main changes

fd0f9a5

github-actions Bot removed the pipelines label Oct 1, 2026

JingyaHuang added 3 commits October 2, 2026 09:52

doc: keep the FLUX.2-dev text encoder off a single TPU chip in the TP…

6ee32b4

… example

doc: encode the prompt on CPU in the TPU TP example so it runs end to…

2d38ce7

… end

test: tighten TPU TP tolerance and simplify the TPU TP worker

1eddd6f

JingyaHuang requested a review from sayakpaul October 2, 2026 12:33

sayakpaul added this to the Release 0.42.0 milestone Oct 5, 2026

sayakpaul reviewed Oct 5, 2026

View reviewed changes

JingyaHuang and others added 8 commits October 5, 2026 13:52

doc: use enable_model_cpu_offload() in the TPU eager example

1351263

doc: shard both FLUX.2-dev text encoder and transformer in the TPU TP…

a170a50

… example

doc: run the text encoders on TPU in the compiled example

5c9e426

test: drop the redundant TPU sync in the TP worker

a3f9ab1

test: trim comments in the TPU TP tests

e7e6c49

test: run the TPU TP test with mp.spawn like the CUDA one

5c262c7

test: pin the TPU TP test to 4 chips

ea2eaca

Merge branch 'main' into add-torchtpu-support

866e165

JingyaHuang requested a review from sayakpaul October 6, 2026 13:48

stevhliu reviewed Oct 6, 2026

View reviewed changes

sayakpaul reviewed Oct 7, 2026

View reviewed changes

sayakpaul and others added 2 commits October 7, 2026 11:55

Merge branch 'main' into add-torchtpu-support

6c5c580

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

2de295b

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

		pipe.transformer = torch.compile(pipe.transformer, backend="tpu", fullgraph=True, dynamic=False)
		pipe.vae = torch.compile(pipe.vae, backend="tpu", fullgraph=True, dynamic=False)

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!

sayakpaul left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JingyaHuang commented Jun 22, 2026 •

edited

Loading

sayakpaul left a comment •

edited

Loading