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

docs/source/en/_toctree.yml

@@ -124,6 +124,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/api/parallel.md

@@ -27,4 +27,4 @@ Parallelism strategies help speed up diffusion transformers by distributing comp
 
 [[autodoc]] TensorParallelConfig
 
-[[autodoc]] hooks.apply_tensor_parallel
+[[autodoc]] hooks.apply_tensor_parallel

docs/source/en/optimization/tpu.md

@@ -0,0 +1,145 @@
+<!--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` | `import torch_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/). After installation,
+`import torch_tpu` registers the `"tpu"` device automatically.
+
+## Eager mode
+
+```python
+import gc
+import torch
+import torch_tpu  # noqa: F401
+
+from diffusers import FluxPipeline
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+
+# 1. Encode on TPU.
+pipe.text_encoder.to("tpu")
+pipe.text_encoder_2.to("tpu")
+with torch.no_grad():
+    prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(
+        prompt="a golden retriever surfing a wave, photorealistic",
+        prompt_2="a golden retriever surfing a wave, photorealistic",
+        device=torch.device("tpu"),
+        max_sequence_length=512,
+    )
+
+# 2. Free the text encoders — nothing below needs them.
+pipe.text_encoder = None
+pipe.text_encoder_2 = None
+gc.collect()
+
+# 3. Move the transformer and VAE in, then denoise with the precomputed embeddings.
+pipe.transformer.to("tpu")
+pipe.vae.to("tpu")
+image = pipe(
+    prompt_embeds=prompt_embeds,
+    pooled_prompt_embeds=pooled_prompt_embeds,
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=0.0,
+).images[0]
+
+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),
+shard it across multiple chips with [`~diffusers.hooks.tensor_parallel.apply_tensor_parallel`], the
+same mechanism [`~ModelMixin.enable_parallelism`] uses for the transformer (see [Tensor
+parallelism](../training/distributed_inference#tensor-parallelism)). It only requires `model:
+torch.nn.Module`, so it works directly on a `transformers.PreTrainedModel` text encoder too, not
+just a diffusers `ModelMixin`. The text encoder doesn't define a `_tp_plan`, so supply one: pair
+each attention/MLP projection that expands the hidden dimension (`"colwise"`) with the one that
+contracts it back (`"rowwise"`), matching the `transformers` model's actual module names.
+
+## Compiled mode
+
+`import torch_tpu` registers `"tpu"` as a `torch.compile` backend name (`TpuBackend` under the hood), so
+components compile like any other `torch.compile` target — no diffusers-specific method needed. The first
+call (warmup) is slow because it compiles; later calls with the same shapes reuse the compiled graph.
+
+> [!IMPORTANT]
+> 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.
+
+```python
+import torch
+import torch_tpu  # noqa: F401 — registers the "tpu" torch.compile backend
+
+from diffusers import FluxPipeline
+
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16,
+)
+pipe.transformer.to("tpu")
+pipe.vae.to("tpu")
+
+pipe.transformer = torch.compile(pipe.transformer, backend="tpu", fullgraph=True, dynamic=False)
+pipe.vae = torch.compile(pipe.vae, backend="tpu", fullgraph=True, dynamic=False)
+
+# Warmup — triggers static graph compilation.
+with torch.no_grad():
+    pipe(
+        prompt="warmup",
+        height=1024,
+        width=1024,
+        num_inference_steps=4,
+        guidance_scale=0.0,
+    )
+
+# Timed inference reuses the compiled graph.
+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")
+```
+
+## Tensor parallelism
+
+Shard a transformer too large for one chip across several with [`~ModelMixin.enable_parallelism`]. Pass a `TensorParallelConfig` with a TPU `DeviceMesh`. For general TP details, (`_tp_plan`, colwise/rowwise), see the [Tensor parallelism](../training/distributed_inference#tensor-parallelism) guide. Set `backend="tpu_dist"` and `DeviceMesh("tpu", ...)` here to enable tensor parallelism.
+
+```python
+import torch
+import torch.distributed as dist
+import torch_tpu  # noqa: F401
+from torch.distributed.device_mesh import DeviceMesh
+
+from diffusers import DiffusionPipeline, TensorParallelConfig
+
+dist.init_process_group(backend="tpu_dist")
+tp_mesh = DeviceMesh("tpu", list(range(dist.get_world_size())))
+
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.2-dev", torch_dtype=torch.bfloat16)
+pipe.transformer.enable_parallelism(config=TensorParallelConfig(mesh=tp_mesh))
+pipe.transformer.to("tpu")
+```

src/diffusers/hooks/tensor_parallel.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import NamedTuple
+
 import torch
 
 from ..models._modeling_parallel import TensorParallelConfig
@@ -20,7 +22,7 @@
 
 logger = get_logger(__name__)  # pylint: disable=invalid-name
 
-_SUPPORTED_TP_DEVICES = ("cuda", "neuron")
+_SUPPORTED_TP_DEVICES = ("cuda", "neuron", "tpu")
 
 
 class PackedColwiseParallel:
@@ -65,6 +67,89 @@ def _blocks_to_block_sizes(total_size: int, blocks: "list[int]") -> "list[int]":
     return [b * unit for b in blocks]
 
 
+class TPShardSpec(NamedTuple):
+    """How one parameter is laid out across the tensor-parallel ranks.
+
+    `dim` is the dimension sharded across ranks, or `None` when the parameter is replicated on every rank (a rowwise
+    bias, which is added after the all-reduce). `block_sizes` partitions `dim` into independently sharded blocks; a
+    plain `"colwise"` / `"rowwise"` style has a single block covering the whole dimension, and packed styles have one
+    per fused projection.
+    """
+
+    dim: "int | None"
+    block_sizes: "list[int] | None"
+
+
+def _local_shard(tensor: torch.Tensor, dim: int, block_sizes: "list[int]", tp_mesh) -> torch.Tensor:
+    """Extract this rank's slice of `tensor` along `dim`.
+
+    Each block is sliced independently and the pieces concatenated, so every fused projection of a packed weight
+    contributes its own contiguous chunk to the rank's shard. A `dim`-1 slice comes back strided, hence the final
+    `contiguous()` — `DTensor.from_local` needs a contiguous local tensor.
+    """
+    rank = tp_mesh.get_local_rank()
+    tp_size = tp_mesh.size()
+
+    parts, offset = [], 0
+    for block_size in block_sizes:
+        # An uneven split is rejected rather than handed to `Shard`, which pads the tail and would break both the
+        # paired colwise/rowwise matmul and the attention head split.
+        if block_size % tp_size != 0:
+            raise ValueError(
+                f"Cannot shard a block of size {block_size} across {tp_size} tensor-parallel ranks: "
+                f"{block_size} is not divisible by {tp_size}."
+            )
+        chunk = block_size // tp_size
+        index = [slice(None)] * tensor.dim()
+        index[dim] = slice(offset + rank * chunk, offset + (rank + 1) * chunk)
+        parts.append(tensor[tuple(index)])
+        offset += block_size
+
+    local = parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim)
+    return local.contiguous()
+
+
+def _block_shard_specs(
+    block: torch.nn.Module, relative_plan: dict
+) -> "list[tuple[torch.nn.Module, str, TPShardSpec]]":
+    """Resolve one block's plan to `(module, param_name, spec)` triples, covering both `weight` and `bias`.
+
+    Modules are returned directly rather than by name so the caller can place the shards without a second lookup.
+    """
+    resolved = []
+    for relative_path, style in relative_plan.items():
+        submodule = block
+        for atom in relative_path.split("."):
+            submodule = getattr(submodule, atom)
+
+        # `_tp_packed_*_blocks` hold absolute sizes rather than proportions; that works because they sum to the
+        # full dimension, so `_blocks_to_block_sizes` computes `unit == 1`.
+        if style == "colwise":
+            weight_spec = TPShardSpec(0, [submodule.weight.shape[0]])
+            bias_spec = weight_spec
+        elif style == "rowwise":
+            weight_spec = TPShardSpec(1, [submodule.weight.shape[1]])
+            bias_spec = TPShardSpec(None, None)
+        elif isinstance(style, PackedColwiseParallel):
+            blocks = style.blocks if style.blocks is not None else submodule._tp_packed_col_blocks
+            weight_spec = TPShardSpec(0, _blocks_to_block_sizes(submodule.weight.shape[0], blocks))
+            bias_spec = weight_spec
+        elif isinstance(style, PackedRowwiseParallel):
+            blocks = style.blocks if style.blocks is not None else submodule._tp_packed_row_blocks
+            weight_spec = TPShardSpec(1, _blocks_to_block_sizes(submodule.weight.shape[1], blocks))
+            bias_spec = TPShardSpec(None, None)
+        else:
+            raise ValueError(
+                f"Unsupported tensor-parallel style '{style}' for '{relative_path}'. "
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+            )
+
+        resolved.append((submodule, "weight", weight_spec))
+        if submodule.bias is not None:
+            resolved.append((submodule, "bias", bias_spec))
+    return resolved
+
+
 def _resolve_tp_plan(model: torch.nn.Module, tp_plan: dict) -> list:
     """Group a flat `_tp_plan` into per-block `(submodule, {relative_path: style})` plans.
 
@@ -193,7 +278,7 @@ def _partition_linear_fn(self, name, module, device_mesh):
     # `distribute_tensor` accepts an indivisible shard dim and just gives the trailing ranks a smaller (or empty)
     # slice, so an uneven split does not raise here — it surfaces much later as a shape or numerics error, because
     # the attention head split and the paired colwise/rowwise Linear both assume equal shards. Reject it up front,
-    # matching what the packed styles above and the Neuron pre-shard path already do.
+    # matching what `_local_shard` already does for the packed styles.
     def _make_checked_col(path: str) -> ColwiseParallel:
         class _CheckedColwiseImpl(ColwiseParallel):
             def _partition_linear_fn(self, name, module, device_mesh):
@@ -240,6 +325,78 @@ def _partition_linear_fn(self, name, module, device_mesh):
     return resolved
 
 
+def _hooks_only_styles(relative_plan: dict) -> dict:
+    """Map a `{relative_path: style}` plan to styles that partition nothing.
+
+    Used when the caller has already placed every planned parameter as a `DTensor`. `parallelize_module` then runs only
+    to register the forward input/output hooks; `_partition_linear_fn` must not re-partition. Packed and plain styles
+    share hook behaviour, so both collapse onto the two styles here.
+
+    Note this is not purely additive: `distribute_module` still replicates any *remaining* plain parameter of the
+    targeted module into a `Replicate()` DTensor via a broadcast, so callers must place every planned parameter
+    themselves.
+    """
+    from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
+
+    class _NoPartitionColwise(ColwiseParallel):
+        def _partition_linear_fn(self, name, module, device_mesh):
+            pass  # weight already Shard(0)
+
+    class _NoPartitionRowwise(RowwiseParallel):
+        def _partition_linear_fn(self, name, module, device_mesh):
+            pass  # weight already Shard(1)
+
+    resolved = {}
+    for path, style in relative_plan.items():
+        if style == "colwise" or isinstance(style, PackedColwiseParallel):
+            resolved[path] = _NoPartitionColwise()
+        elif style == "rowwise" or isinstance(style, PackedRowwiseParallel):
+            resolved[path] = _NoPartitionRowwise()
+        else:
+            raise ValueError(
+                f"Unsupported tensor-parallel style '{style}' for '{path}'. "
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+            )
+    return resolved
+
+
+def _pre_shard_and_parallelize(tp_mesh, groups: list, device: torch.device) -> None:
+    """Slice every planned parameter on CPU and move only this rank's shard to `device`.
+
+    The default path lets `parallelize_module` distribute the weights, which materializes each full weight on every
+    rank before scattering it, so peak memory per rank is the size of the whole weight even though only a shard
+    survives. Slicing first and handing `DTensor.from_local` just this rank's piece keeps the full tensor off the
+    accelerator, which is what allows sharding a model larger than one device's memory. `parallelize_module` then runs
+    over `_hooks_only_styles`, distributing nothing and registering only the hooks the forward pass needs.
+
+    `device` is the only backend-specific input; the slicing itself is the same on any accelerator. Unlike the default
+    path this does not broadcast from a single rank, so every rank must already hold the same weights — true after
+    loading a checkpoint, not after a random init.
+
+    Model weights must be on CPU when this is called.
+    """
+    import torch.nn as nn
+    from torch.distributed.tensor import DTensor, Replicate, Shard
+    from torch.distributed.tensor.parallel import parallelize_module
+
+    for block, relative_plan in groups:
+        for submodule, param_name, spec in _block_shard_specs(block, relative_plan):
+            param = getattr(submodule, param_name)
+            if spec.dim is None:
+                # A rowwise bias is added after the all-reduce, so every rank needs the whole vector.
+                local, placement = param.data, Replicate()
+            else:
+                local, placement = _local_shard(param.data, spec.dim, spec.block_sizes, tp_mesh), Shard(spec.dim)
+            submodule.register_parameter(
+                param_name,
+                nn.Parameter(
+                    DTensor.from_local(local.to(device), tp_mesh, [placement]),
+                    requires_grad=param.requires_grad,
+                ),
+            )
+        parallelize_module(block, tp_mesh, _hooks_only_styles(relative_plan))
+
+
 def apply_tensor_parallel(
     model: torch.nn.Module,
     config: TensorParallelConfig,
@@ -257,7 +414,12 @@ def apply_tensor_parallel(
             f"or from the active accelerator when the mesh is built from `tp_degree`."
         )
 
-    backend = "neuron" if tp_mesh.device_type == "neuron" else "default"
+    if tp_mesh.device_type == "neuron":
+        backend = "neuron"
+    elif tp_mesh.device_type == "tpu":
+        backend = "tpu"
+    else:
+        backend = "default"
     groups = _resolve_tp_plan(model, tp_plan)
     logger.debug(f"Applying tensor parallel (backend={backend}) over {len(groups)} module group(s) on mesh {tp_mesh}.")
 
@@ -267,6 +429,14 @@ def apply_tensor_parallel(
         _apply_tp_neuron(model, tp_mesh, groups)
         return
 
+    if backend == "tpu":
+        # Pre-shard rather than let `parallelize_module` distribute: it materializes each full weight on every chip
+        # before scattering it, which exhausts HBM for a large diffusion transformer. Address the chip as "tpu" with
+        # no index — "tpu:rank" would mean chip `rank` from this process's view, but each torchrun worker only has
+        # access to its own assigned chip.
+        _pre_shard_and_parallelize(tp_mesh, groups, torch.device("tpu"))
+        return
+
     from torch.distributed.tensor.parallel import parallelize_module
 
     for submodule, relative_plan in groups:

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`):