# Google TPU

张量处理单元（TPU）是 Google 的机器学习加速器，以云 TPU（v6e、v7x）的形式提供。Diffusers 通过 [TorchTPU](https://github.com/google-pytorch/torch_tpu/)（TPU 的 PyTorch 后端）在其上运行，只需对代码进行极少量的修改。

共有三种执行模式可用：

| 模式 | 常量 | 如何激活 | 备注 |
|---|---|---|---|
| 严格 eager（默认） | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | 一次分派一个操作，异步 |
| 优化的 eager | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | 操作被延迟并融合到更大的计算图中；比严格即时执行更快，但编译时间更长 |
| 编译 | — | `torch.compile(module, backend="tpu")` | 使用 `TpuBackend` 进行 AOT 编译 |

请按照 [TorchTPU 安装指南](https://github.com/google-pytorch/torch_tpu/) 进行操作。安装完成后，`import torch` 会自动加载并注册 `"tpu"` 设备，因此只需将 `pipe.to("tpu")` 作为唯一改动即可。仅当通过 `TORCH_DEVICE_BACKEND_AUTOLOAD=0` 禁用了后端自动加载时，才需要额外添加 `import torch_tpu`。

## Eager 模式

FLUX.1-schnell 无法一次性完整加载到单个 v6e 芯片上，因此请使用 [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload)，仅在模型运行时将其移至 TPU。该函数会自动检测 `"tpu"` 设备。

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

如果模型对于单个芯片来说过大，或者您拥有多个芯片并希望降低延迟，请将模型分片到多个芯片上。请参阅 [张量并行](#tensor-parallelism) 部分。

## 编译模式

TorchTPU 将 `"tpu"` 注册为 `torch.compile` 的后端名称（底层为 `TpuBackend`），因此组件的编译方式与其他 `torch.compile` 目标相同。首次调用（预热）速度较慢，因为需要进行编译；后续使用相同形状的调用会复用已编译的图。

> [!IMPORTANT]
> TorchTPU 需要**静态形状**，因此请传入 `dynamic=False`。新的 `height` 或 `width` 会针对该形状重新编译一次；已见过的形状会被复用。更改 `num_inference_steps` 不会触发重新编译。

当整个流水线可以容纳在单个芯片上时，将其迁移到 TPU 并编译完整的 transformer。Stable Diffusion 3.5 Medium（在 bf16 下约 15GB）可以容纳在单个 v6e 芯片上。

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

如果流水线无法在单个芯片上运行：

- 如果使用多块芯片，请改用 [张量并行](#tensor-parallelism) 进行分片。所有内容都保留在 TPU 上，并且 `pipe.transformer.compile(...)` 在分片后的 transformer 上的工作方式相同。
- 使用 [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) 时，offload 钩子无法被 `torch.compile` 追踪。请仅编译 transformer 的重复块，使用 `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`。

## 张量并行

将超出单芯片容量的模型分片到多个芯片上。FLUX.2-dev 的文本编码器（约 48GB）和 Transformer（约 64GB）均超出单芯片容量，因此以下示例对两者都进行了分片：

- 使用 [TensorParallelConfig](/docs/diffusers/v0.43.0/zh/api/parallel#diffusers.TensorParallelConfig) 配置 transformer，并将其传递给 [from_pretrained()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.from_pretrained) 的 `parallel_config` 参数。每个 rank 仅读取每个分片权重中属于自己的切片，因此完整的模型永远不会被实例化。有关 TP 的通用细节（`_tp_plan`、colwise/rowwise），请参阅 [Tensor parallelism](./tensor_parallelism) 指南。
- 使用 Transformers 自身的 [张量并行](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) 来配置文本编码器，传入 [DistributedConfig](https://huggingface.co/docs/transformers/v5.22.2/en/expert_parallelism#transformers.DistributedConfig) 以及相同的 mesh。

在 TPU 上，使用 `backend="tpu_dist"` 初始化进程组，并使用 `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")
```

为每个芯片启动一个进程。将 `--nproc_per_node` 设置为使用主机上所有 TPU 芯片的数量。

```bash
eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /')
torchrun --nproc_per_node=8 flux2_tp.py
```

> [!WARNING]
> 上述启动命令使用了 TorchTPU 的内部模块（`torch_tpu._internal`），而非稳定的 API，预计将在未来的 TorchTPU 版本中发生变更。请查阅 [TorchTPU 仓库](https://github.com/google-pytorch/torch_tpu/) 以了解当前启动多芯片任务的方式。