Google TPU

Les Tensor Processing Units (TPUs) sont les accélérateurs d’apprentissage automatique de Google, disponibles sous forme de Cloud TPUs (v6e, v7x). Diffusers s’exécute dessus via TorchTPU, le backend PyTorch pour les TPUs, avec des modifications de code minimales.

Trois modes d'exécution sont disponibles :

Suivez le guide d’installation de TorchTPU. Une fois installé, import torch le charge automatiquement et enregistre le périphérique "tpu", de sorte que pipe.to("tpu") est la seule modification nécessaire. Ajoutez import torch_tpu uniquement si vous avez désactivé le chargement automatique du backend avec TORCH_DEVICE_BACKEND_AUTOLOAD=0.

Mode Eager

FLUX.1-schnell ne tient pas entièrement sur une seule puce v6e, utilisez donc enable_model_cpu_offload() pour déplacer chaque modèle sur le TPU uniquement pendant son exécution. Il détecte automatiquement le périphérique "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 modèle est trop volumineux pour une puce unique, ou si vous disposez de plusieurs puces et souhaitez réduire la latence, répartissez les modèles sur plusieurs puces. Consultez la section Parallélisme de tenseurs.

Mode compilé

TorchTPU enregistre "tpu" comme nom de backend pour torch.compile (TpuBackend en interne), de sorte que les composants sont compilés comme tout autre objectif torch.compile. Le premier appel (phase de préchauffage) est lent car il effectue la compilation ; les appels ultérieurs avec les mêmes formes réutilisent le graphe compilé.

TorchTPU exige des formes statiques, il faut donc passer dynamic=False. Une nouvelle height ou width déclenche une nouvelle compilation pour cette forme, une seule fois ; les formes déjà vues sont réutilisées. Modifier num_inference_steps ne provoque pas de re-compilation.

Lorsque le pipeline complet tient sur une seule puce, déplacez-le sur le TPU et compilez le transformer complet. Stable Diffusion 3.5 Medium (~15GB en bf16) tient sur une seule puce 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 le pipeline ne tient pas sur une seule puce :

• Avec plusieurs puces, utilisez plutôt le parallélisme de tenseurs pour le sharding. Tout reste sur le TPU, et pipe.transformer.compile(...) fonctionne de la même manière sur le transformer sharded.
• Avec enable_model_cpu_offload(), les hooks de déchargement ne peuvent pas être tracés par torch.compile. Compilez uniquement les blocs répétés du transformeur, avec pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False).

Parallélisme de tenseurs

Répartissez les modèles trop volumineux pour une seule puce sur plusieurs. L’encodeur de texte (~48 Go) et le transformeur (~64 Go) de FLUX.2-dev dépassent chacun la capacité d’une puce unique, l’exemple ci-dessous répartit donc les deux :

• le transformer avec TensorParallelConfig, transmis à l’argument parallel_config de from_pretrained(). Chaque rang ne lit que sa propre tranche de chaque poids fractionné, de sorte que le modèle complet n’est jamais matérialisé. Pour les détails généraux sur le TP (_tp_plan, colwise/rowwise), consultez le guide Tensor parallelism.
• l’encodeur de texte avec le parallélisme de tenseurs propre à Transformers, en passant un DistributedConfig et le même maillage.

Sur TPU, initialisez le groupe de processus avec backend="tpu_dist" et construisez le maillage avec 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")

Lancez un processus par puce. Définissez --nproc_per_node pour utiliser le nombre total de puces TPU sur votre hôte.

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

La commande de lancement ci-dessus utilise un module interne de TorchTPU (torch_tpu._internal), et non une API stable, et est susceptible d'évoluer dans une future version de TorchTPU. Consultez le dépôt TorchTPU pour connaître la méthode actuelle de lancement des tâches multi-puces.