Skip to content

[TPU] TorchTPU backend integration - eager / torch.compile / tp - #14039

Merged
sayakpaul merged 79 commits into
huggingface:mainfrom
JingyaHuang:add-torchtpu-support
Oct 9, 2026
Merged

sayakpaul merged 79 commits into
huggingface:mainfrom
JingyaHuang:add-torchtpu-support

Conversation

@JingyaHuang

@JingyaHuang JingyaHuang commented Jun 22, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Need the fix #14739 and perferrably merge the sharding improvement PR #14544 first.

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.

@github-actions 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 JingyaHuang changed the title [TPU] Initial TorchTPU backend integration (eager + torch.compile) [TPU] TorchTPU backend integration - eager / torch.compile / tp Sep 7, 2026
@JingyaHuang
JingyaHuang marked this pull request as ready for review September 7, 2026 16:23
…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>
@JingyaHuang
JingyaHuang requested a review from sayakpaul October 6, 2026 13:48
Comment thread docs/source/en/optimization/tpu.md Outdated
Comment thread docs/source/en/optimization/tpu.md Outdated
Comment thread docs/source/en/optimization/tpu.md Outdated

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! I left some further comments.

Comment thread docs/source/en/optimization/tpu.md
Comment thread tests/models/testing_utils/parallelism.py
Comment thread tests/models/testing_utils/parallelism.py Outdated
@JingyaHuang
JingyaHuang requested a review from sayakpaul October 7, 2026 16:36

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Just one major comment regarding the SD3 change.

Comment thread src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py Outdated
Comment thread tests/models/testing_utils/parallelism.py
Comment thread docs/source/en/optimization/tpu.md Outdated
Comment thread docs/source/en/optimization/tpu.md
@JingyaHuang
JingyaHuang requested a review from sayakpaul October 9, 2026 09:07
@sayakpaul
sayakpaul merged commit 1d5d056 into huggingface:main Oct 9, 2026
27 of 29 checks passed
@JingyaHuang
JingyaHuang deleted the add-torchtpu-support branch October 9, 2026 12:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation hooks models pipelines size/L PR with diff > 200 LOC tests utils

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants