Google TPU

Las unidades de procesamiento de tensores (TPU) son los aceleradores de aprendizaje automático de Google, disponibles como Cloud TPU (v6e, v7x). Diffusers se ejecuta en ellos a través de TorchTPU, el backend de PyTorch para TPU, con cambios mínimos en el código.

Hay disponibles tres modos de ejecución:

Siga la guía de instalación de TorchTPU. Una vez instalado, import torch lo carga automáticamente y registra el dispositivo "tpu", por lo que pipe.to("tpu") es el único cambio necesario. Añada import torch_tpu solo si deshabilitó la carga automática del backend con TORCH_DEVICE_BACKEND_AUTOLOAD=0.

Modo eager

FLUX.1-schnell no cabe en un solo chip v6e a la vez, por lo que use enable_model_cpu_offload() para mover cada modelo al TPU solo mientras se ejecuta. Detecta automáticamente el 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")

Si un modelo es demasiado grande para un solo chip, o si tienes varios chips y quieres una latencia más baja, divide los modelos entre los chips. Consulta la sección Paralelismo de tensores.

Modo compilado

TorchTPU registra "tpu" como nombre de backend de torch.compile (TpuBackend por debajo), por lo que los componentes se compilan como cualquier otro destino de torch.compile. La primera llamada (calentamiento) es lenta porque compila; las llamadas posteriores con las mismas formas reutilizan el gráfico compilado.

TorchTPU requiere formas estáticas, por lo que debe pasar dynamic=False. Una nueva height o width vuelve a compilar para esa forma, una sola vez; las formas ya vistas se reutilizan. Cambiar num_inference_steps no vuelve a compilar.

Cuando el pipeline completo cabe en un solo chip, muévelo a la TPU y compila el transformador completo. Stable Diffusion 3.5 Medium (~15GB en bf16) cabe en un solo 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")

Si el pipeline no cabe en un solo chip:

• Con varios chips, divídalo en fragmentos mediante paralelismo de tensores en su lugar. Todo permanece en la TPU, y pipe.transformer.compile(...) funciona de la misma manera en el transformador fragmentado.
• Con enable_model_cpu_offload(), los hooks de offload no pueden ser trazados por torch.compile. En su lugar, compile solo los bloques repetidos del transformer, con pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False).

Paralelismo de tensores

Divida los modelos demasiado grandes para un solo chip en varios. El codificador de texto (~48GB) y el transformador (~64GB) de FLUX.2-dev superan cada uno la capacidad de un solo chip, por lo que el ejemplo siguiente divide ambos:

• el transformer con TensorParallelConfig, pasado al argumento parallel_config de from_pretrained(). Cada rank solo lee su propia porción de cada peso fragmentado, por lo que el modelo completo nunca se materializa. Para más detalles generales sobre TP (_tp_plan, colwise/rowwise), consulte la guía de Tensor parallelism.
• el codificador de texto con el propio paralelismo de tensores de Transformers, pasando un DistributedConfig y la misma malla.

En TPU, inicializa el grupo de procesos con backend="tpu_dist" y construye la malla con 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 un proceso por chip. Establezca --nproc_per_node para usar el número total de chips TPU de su host.

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

El comando de lanzamiento anterior utiliza un módulo interno de TorchTPU (torch_tpu._internal), no una API estable, y se espera que cambie en una futura versión de TorchTPU. Consulte el repositorio de TorchTPU para conocer la forma actual de lanzar trabajos de varios chips.