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

docs/source/en/_toctree.yml

@@ -114,6 +114,8 @@
       title: Intel Gaudi
     - local: optimization/neuron
       title: AWS Neuron
+    - local: optimization/tpu
+      title: TPU
     title: Hardware-specific acceleration
   - isExpanded: false
     sections:

docs/source/en/optimization/tpu.md

@@ -0,0 +1,151 @@
+<!--Copyright 2026 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# TorchTPU
+
+[TorchTPU](https://github.com/google-pytorch/torch_tpu/) is a PyTorch backend for Google's Tensor Processing Units (TPUs), which lets you run Diffusers pipelines on Cloud TPUs (v6e, v5p, etc.) with minimal code changes.
+
+Two execution modes are available:
+
+| Mode | Constant | How to activate | Notes |
+|---|---|---|---|
+| Strict eager (default) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Operations dispatched one at a time, asynchronous |
+| Compile | — | `torch.compile(module, backend="tpu")` | AOT compilation with `TpuBackend` |
+
+Follow the [TorchTPU installation guide](https://github.com/google-pytorch/torch_tpu/). Once installed, `import torch`
+loads it automatically and registers the `"tpu"` device, so `pipe.to("tpu")` is the only change needed. Add
+`import torch_tpu` only if you disabled backend autoloading with `TORCH_DEVICE_BACKEND_AUTOLOAD=0`.
+
+## Eager mode
+
+FLUX.1-schnell doesn't fit on a single v6e chip all at once, so use [`~DiffusionPipeline.enable_model_cpu_offload`] to
+move each model to the TPU only while it runs. It detects the `"tpu"` device automatically.
+
+```python
+import torch
+
+from diffusers import FluxPipeline
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload()
+
+image = pipe(
+    prompt="a golden retriever surfing a wave, photorealistic",
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=0.0,
+).images[0]
+
+image.save("output.png")
+```
+
+If a model is too large for a single chip, or you have several chips and want lower latency, shard the models
+across chips instead. See the [Tensor parallelism](#tensor-parallelism) section.
+
+## Compiled mode
+
+TorchTPU registers `"tpu"` as a `torch.compile` backend name (`TpuBackend` under the hood), so
+components compile like any other `torch.compile` target. The first
+call (warmup) is slow because it compiles; later calls with the same shapes reuse the compiled graph.
+
+> [!IMPORTANT]
+> TorchTPU requires **static shapes**, so pass `dynamic=False`. A new `height` or `width` compiles again for that
+> shape, once; shapes already seen are reused. Changing `num_inference_steps` doesn't recompile.
+
+When the whole pipeline fits on one chip, move it to the TPU and compile the full transformer. Stable Diffusion 3.5
+Medium (~15GB in bf16) fits on a single v6e chip.
+
+```python
+import torch
+
+from diffusers import StableDiffusion3Pipeline
+
+pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16)
+pipe.to("tpu")
+pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
+
+# Warmup — triggers static graph compilation.
+pipe(prompt="warmup", height=1024, width=1024, num_inference_steps=40, guidance_scale=4.5)
+
+# Later calls with the same shapes reuse the compiled graph.
+image = pipe(
+    prompt="a golden retriever surfing a wave, photorealistic",
+    height=1024,
+    width=1024,
+    num_inference_steps=40,
+    guidance_scale=4.5,
+).images[0]
+
+image.save("output.png")
+```
+
+If the pipeline doesn't fit on one chip:
+
+- With several chips, shard it with [tensor parallelism](#tensor-parallelism) instead. Everything stays on the TPU,
+  and `pipe.transformer.compile(...)` works the same way on the sharded transformer.
+- With [`~DiffusionPipeline.enable_model_cpu_offload`], the offload hooks can't be traced by `torch.compile`. Compile
+  only the transformer's repeated blocks instead, with
+  `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`.
+
+## Tensor parallelism
+
+Shard models too large for one chip across several. FLUX.2-dev's text encoder (~48GB) and transformer (~64GB) each
+exceed a single chip, so the example below shards both:
+
+- the transformer with [`TensorParallelConfig`], passed to the `parallel_config` argument of [`~ModelMixin.from_pretrained`]. Each rank reads only its own slice of every sharded weight, so the full model is never materialized. For general TP details (`_tp_plan`, colwise/rowwise), see the [Tensor parallelism](../training/distributed_inference#tensor-parallelism) guide.
+- the text encoder with Transformers' own [tensor parallelism](https://huggingface.co/docs/transformers/perf_infer_gpu_multi), passing a [`~transformers.DistributedConfig`] and the same mesh.
+
+On TPU, initialize the process group with `backend="tpu_dist"` and build the mesh with `DeviceMesh("tpu", ...)`.
+
+```python
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh
+from transformers import DistributedConfig, Mistral3ForConditionalGeneration
+
+from diffusers import Flux2Pipeline, Flux2Transformer2DModel, TensorParallelConfig
+
+dist.init_process_group(backend="tpu_dist")
+mesh = DeviceMesh("tpu", list(range(dist.get_world_size())))
+
+repo_id = "black-forest-labs/FLUX.2-dev"
+text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
+    repo_id,
+    subfolder="text_encoder",
+    dtype=torch.bfloat16,
+    distributed_config=DistributedConfig(tp_plan="auto"),
+    device_mesh=mesh,
+)
+transformer = Flux2Transformer2DModel.from_pretrained(
+    repo_id, subfolder="transformer", dtype=torch.bfloat16, parallel_config=TensorParallelConfig(mesh=mesh)
+)
+pipe = Flux2Pipeline.from_pretrained(
+    repo_id, text_encoder=text_encoder, transformer=transformer, dtype=torch.bfloat16
+)
+pipe.vae.to("tpu")
+
+image = pipe(
+    prompt="a golden retriever surfing a wave, photorealistic",
+    num_inference_steps=28,
+    generator=torch.Generator("cpu").manual_seed(0),
+).images[0]
+if dist.get_rank() == 0:
+    image.save("output.png")
+```
+
+Launch one process per chip. Set `--nproc_per_node` to use all the number of TPU chips on your host.
+
+```bash
+eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /')
+torchrun --nproc_per_node=8 flux2_tp.py
+```

src/diffusers/hooks/tensor_parallel.py

@@ -22,7 +22,7 @@
 
 logger = get_logger(__name__)  # pylint: disable=invalid-name
 
-_SUPPORTED_TP_DEVICES = ("cuda", "neuron")
+_SUPPORTED_TP_DEVICES = ("cuda", "neuron", "tpu")
 
 
 class PackedColwiseParallel:

src/diffusers/models/_modeling_parallel.py

@@ -161,7 +161,7 @@ class TensorParallelConfig:
     Tensor parallelism shards weight matrices (column-wise and row-wise) across devices. Each device computes a partial
     result; an AllReduce/AllGather at layer boundaries reconstructs the full output. Uses
     `torch.distributed.tensor.parallelize_module` with `ColwiseParallel` / `RowwiseParallel` sharding styles. Supported
-    device types are `"cuda"` and `"neuron"`.
+    device types are `"cuda"`, `"neuron"` and `"tpu"`.
 
     Args:
         tp_degree (`int`, defaults to `1`):

src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -1066,7 +1066,8 @@ def __call__(
                 # expand the latents if we are doing classifier free guidance
                 latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0])
+                # `clone()` so it isn't a view into `timesteps`, which would recompile `torch.compile` every step
+                timestep = t.expand(latent_model_input.shape[0]).clone()
 
                 noise_pred = self.transformer(
                     hidden_states=latent_model_input,
@@ -1088,7 +1089,7 @@ def __call__(
                         else False
                     )
                     if skip_guidance_layers is not None and should_skip_layers:
-                        timestep = t.expand(latents.shape[0])
+                        timestep = t.expand(latents.shape[0]).clone()
                         latent_model_input = latents
                         noise_pred_skip_layers = self.transformer(
                             hidden_states=latent_model_input,

src/diffusers/utils/__init__.py

@@ -116,6 +116,7 @@
     is_torch_mlu_available,
     is_torch_neuronx_available,
     is_torch_npu_available,
+    is_torch_tpu_available,
     is_torch_version,
     is_torch_xla_available,
     is_torch_xla_version,

src/diffusers/utils/import_utils.py

@@ -178,6 +178,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b
 _torch_xla_available, _torch_xla_version = _is_package_available("torch_xla")
 _torch_npu_available, _torch_npu_version = _is_package_available("torch_npu")
 _torch_mlu_available, _torch_mlu_version = _is_package_available("torch_mlu")
+_torch_tpu_available, _torch_tpu_version = _is_package_available("torch_tpu")
 _torch_neuronx_available, _torch_neuronx_version = _is_package_available("torch_neuronx")
 _transformers_available, _transformers_version = _is_package_available("transformers")
 _hf_hub_available, _hf_hub_version = _is_package_available("huggingface_hub")
@@ -238,6 +239,10 @@ def is_torch_mlu_available():
     return _torch_mlu_available
 
 
+def is_torch_tpu_available():
+    return _torch_tpu_available
+
+
 def is_torch_neuronx_available():
     return _torch_neuronx_available
 

tests/models/testing_utils/__init__.py

@@ -23,6 +23,7 @@
     ContextParallelAttentionBackendsTesterMixin,
     ContextParallelTesterMixin,
     TensorParallelTesterMixin,
+    TensorParallelTPUTesterMixin,
 )
 from .quantization import (
     AutoRoundCompileTesterMixin,
@@ -67,6 +68,7 @@
     "ContextParallelTesterMixin",
     "ContextParallelAttentionBackendsTesterMixin",
     "TensorParallelTesterMixin",
+    "TensorParallelTPUTesterMixin",
     "CPUOffloadTesterMixin",
     "FasterCacheConfigMixin",
     "FasterCacheTesterMixin",