Google TPU

Tensor Processing Units (TPU) là bộ tăng tốc máy học của Google, có sẵn dưới dạng Cloud TPU (v6e, v7x). Diffusers chạy trên chúng thông qua TorchTPU, backend PyTorch cho TPU, với các thay đổi mã tối thiểu.

Có ba chế độ thực thi:

Làm theo hướng dẫn cài đặt TorchTPU. Sau khi cài đặt, import torch sẽ tự động nạp và đăng ký thiết bị "tpu", do đó pipe.to("tpu") là thay đổi duy nhất cần thực hiện. Chỉ thêm import torch_tpu nếu bạn đã tắt tự động nạp backend bằng TORCH_DEVICE_BACKEND_AUTOLOAD=0.

Chế độ Eager

FLUX.1-schnell không vừa trên một chip v6e duy nhất cùng lúc, vì vậy hãy dùng enable_model_cpu_offload() để chuyển từng model sang TPU chỉ trong khi nó chạy. Nó tự động phát hiện thiết bị "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")

Nếu một model quá lớn cho một chip duy nhất, hoặc bạn có nhiều chip và muốn giảm độ trễ, hãy chia model thành các shard trên các chip thay vì chạy trên một chip. Xem phần Tensor parallelism.

Chế độ biên dịch

TorchTPU đăng ký "tpu" làm tên backend của torch.compile (bên trong là TpuBackend), do đó các thành phần được biên dịch giống như mọi mục tiêu torch.compile khác. Lệnh gọi đầu tiên (giai đoạn khởi động) sẽ chậm vì quá trình biên dịch; các lệnh gọi sau với cùng các hình dạng sẽ tái sử dụng đồ thị đã biên dịch.

TorchTPU yêu cầu hình dạng tĩnh, vì vậy hãy truyền dynamic=False. Một height hoặc width mới sẽ được biên dịch lại cho hình dạng đó, chỉ một lần; các hình dạng đã gặp trước đó sẽ được tái sử dụng. Việc thay đổi num_inference_steps không gây ra biên dịch lại.

Khi toàn bộ pipeline vừa trên một chip, hãy di chuyển nó sang TPU và biên dịch toàn bộ transformer. Stable Diffusion 3.5 Medium (~15GB trong bf16) vừa trên một chip v6e đơn lẻ.

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

Nếu pipeline không vừa trên một chip:

• Với nhiều chip, hãy shard bằng tensor parallelism thay thế. Mọi thứ vẫn nằm trên TPU, và pipe.transformer.compile(...) hoạt động theo cách tương tự trên transformer đã shard.
• Với enable_model_cpu_offload(), các hook offload không thể được theo dõi bởi torch.compile. Hãy chỉ compile các khối lặp lại của transformer, bằng cách sử dụng pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False).

Song song tensor

Chia nhỏ các mô hình quá lớn cho một chip duy nhất thành nhiều chip. Bộ mã hóa văn bản (~48GB) và transformer (~64GB) của FLUX.2-dev đều vượt quá dung lượng của một chip, vì vậy ví dụ dưới đây chia nhỏ cả hai:

• transformer với TensorParallelConfig, được truyền vào đối số parallel_config của from_pretrained(). Mỗi rank chỉ đọc phần riêng của mình trong mỗi trọng số được chia nhỏ, nên mô hình đầy đủ không bao giờ được tạo ra. Để biết chi tiết chung về TP (_tp_plan, colwise/rowwise), hãy xem hướng dẫn Tensor parallelism.
• bộ mã hóa văn bản với song song tensor của riêng Transformers, truyền vào một DistributedConfig và cùng một lưới (mesh).

Trên TPU, khởi tạo nhóm tiến trình với backend="tpu_dist" và xây dựng lưới bằng 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")

Khởi chạy một tiến trình cho mỗi chip. Đặt --nproc_per_node để sử dụng toàn bộ số chip TPU trên máy chủ của bạn.

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

Lệnh khởi chạy ở trên sử dụng một module nội bộ của TorchTPU (torch_tpu._internal), không phải là API ổn định, và dự kiến sẽ thay đổi trong một bản phát hành TorchTPU tương lai. Hãy kiểm tra kho TorchTPU để biết cách khởi chạy các tác vụ đa chip hiện tại.