diffusers/v0.43.0/zh/optimization/tpu.md

Google TPU

张量处理单元（TPU）是 Google 的机器学习加速器，以云 TPU（v6e、v7x）的形式提供。Diffusers 通过 TorchTPU（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 安装指南 进行操作。安装完成后，import torch 会自动加载并注册 "tpu" 设备，因此只需将 pipe.to("tpu") 作为唯一改动即可。仅当通过 TORCH_DEVICE_BACKEND_AUTOLOAD=0 禁用了后端自动加载时，才需要额外添加 import torch_tpu。

Eager 模式

FLUX.1-schnell 无法一次性完整加载到单个 v6e 芯片上，因此请使用 enable_model_cpu_offload()，仅在模型运行时将其移至 TPU。该函数会自动检测 "tpu" 设备。

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

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

编译模式

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

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

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

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

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

• 如果使用多块芯片，请改用 张量并行 进行分片。所有内容都保留在 TPU 上，并且 pipe.transformer.compile(...) 在分片后的 transformer 上的工作方式相同。
• 使用 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 配置 transformer，并将其传递给 from_pretrained() 的 parallel_config 参数。每个 rank 仅读取每个分片权重中属于自己的切片，因此完整的模型永远不会被实例化。有关 TP 的通用细节（_tp_plan、colwise/rowwise），请参阅 Tensor parallelism 指南。
• 使用 Transformers 自身的 张量并行 来配置文本编码器，传入 DistributedConfig 以及相同的 mesh。

在 TPU 上，使用 backend="tpu_dist" 初始化进程组，并使用 DeviceMesh("tpu", ...) 构建网格。

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 芯片的数量。

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

上述启动命令使用了 TorchTPU 的内部模块（torch_tpu._internal），而非稳定的 API，预计将在未来的 TorchTPU 版本中发生变更。请查阅 TorchTPU 仓库 以了解当前启动多芯片任务的方式。