Google TPU

Tensor Processing Units (TPUs) são os aceleradores de machine learning da Google, disponíveis como Cloud TPUs (v6e, v7x). O Diffusers roda neles por meio do TorchTPU, o backend do PyTorch para TPUs, com mudanças mínimas no código.

Três modos de execução estão disponíveis:

Siga o guia de instalação do TorchTPU. Uma vez instalado, import torch o carrega automaticamente e registra o dispositivo "tpu", de modo que pipe.to("tpu") é a única alteração necessária. Adicione import torch_tpu apenas se você desativou o carregamento automático do backend com TORCH_DEVICE_BACKEND_AUTOLOAD=0.

Modo Eager

FLUX.1-schnell não cabe em um único chip v6e de uma vez, então use enable_model_cpu_offload() para mover cada modelo para o TPU apenas enquanto ele está em execução. Ele detecta automaticamente o dispositivo "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")

Se um modelo for grande demais para um único chip, ou se você tiver vários chips e quiser menor latência, distribua os modelos entre os chips. Consulte a seção Paralelismo de tensor.

Modo compilado

TorchTPU registra "tpu" como um nome de backend para torch.compile (TpuBackend por baixo dos panos), de modo que os componentes são compilados como qualquer outro alvo de torch.compile. A primeira chamada (aquecimento) é lenta porque realiza a compilação; chamadas subsequentes com as mesmas formas reutilizam o gráfico compilado.

TorchTPU requer formas estáticas, então passe dynamic=False. Um novo height ou width compila novamente para essa forma, uma vez; formas já vistas são reutilizadas. Alterar num_inference_steps não recompila.

Quando o pipeline inteiro cabe em um único chip, mova-o para o TPU e compile o transformer completo. Stable Diffusion 3.5 Medium (~15GB em bf16) cabe em um único chip 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")

Se o pipeline não couber em um único chip:

• Com vários chips, use tensor parallelism para fazer o sharding. Tudo permanece no TPU, e pipe.transformer.compile(...) funciona da mesma forma no transformer shardado.
• Com enable_model_cpu_offload(), os hooks de offload não podem ser rastreados pelo torch.compile. Compile apenas os blocos repetidos do transformer, com pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False).

Paralelismo de tensor

Divida modelos grandes demais para um único chip entre vários. O codificador de texto (~48GB) e o transformer (~64GB) do FLUX.2-dev excedem, cada um, a capacidade de um único chip, portanto o exemplo abaixo divide ambos:

• o transformer com TensorParallelConfig, passado para o argumento parallel_config de from_pretrained(). Cada rank lê apenas a própria fatia de cada peso particionado, de modo que o modelo completo nunca é materializado. Para detalhes gerais de TP (_tp_plan, colwise/rowwise), consulte o guia Tensor parallelism.
• o codificador de texto com o próprio paralelismo de tensor dos Transformers, passando um DistributedConfig e a mesma malha.

No TPU, inicialize o grupo de processos com backend="tpu_dist" e construa a malha com 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")

Inicie um processo por chip. Defina --nproc_per_node para usar o número total de chips TPU no seu host.

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

O comando de inicialização acima usa um módulo interno do TorchTPU (torch_tpu._internal), não uma API estável, e é esperado que mude em uma futura liberação do TorchTPU. Consulte o repositório do TorchTPU para a maneira atual de iniciar trabalhos multi-chip.