Hugging Face's logo

Diffusers documentation

Google TPU

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Google TPU

张量处理单元（TPU）是 Google 的机器学习加速器，以云 TPU（v6e、v7x）的形式提供。 Diffusers 通过 TorchTPU（TPU 的 PyTorch 后端）在其上运行， 只需对代码进行极少的修改。

共有三种执行模式可用：

模式常量如何激活备注
严格 eager（默认）EagerMode.DEFER_NEVERpipe.to("tpu")一次分派一个操作，异步
优化的 eagerEagerMode.DEFER_AND_FUSETORCH_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 仓库 以了解当前启动多芯片任务的方式。

Update on GitHub

←AWS Trainium CacheDiT→