# 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](https://github.com/google-pytorch/torch_tpu/), le backend PyTorch pour les TPUs, avec des modifications de code minimales.

Trois modes d'exécution sont disponibles :

| Mode | Constant | Comment activer | Notes |
|---|---|---|---|
| Strict eager (par défaut) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Opérations expédiées une à la fois, asynchrones |
| Eager optimisé | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | Opérations différées et fusionnées en graphes plus larges ; plus rapide que l'exécution stricte en mode eager, compilation plus longue |
| Compiler | — | `torch.compile(module, backend="tpu")` | Compilation AOT avec `TpuBackend` |

Suivez le [guide d’installation de TorchTPU](https://github.com/google-pytorch/torch_tpu/). 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()](/docs/diffusers/main/fr/api/pipelines/overview#diffusers.DiffusionPipeline.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"`.

```python
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 [Tensor parallelism](#tensor-parallelism).

## 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 de `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é.

> [!IMPORTANT]
> 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 entier. Stable Diffusion 3.5 Medium (~15 Go en bf16) tient sur une seule puce v6e.

```python
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](#tensor-parallelism) 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()](/docs/diffusers/main/fr/api/pipelines/overview#diffusers.DiffusionPipeline.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](/docs/diffusers/main/fr/api/parallel#diffusers.TensorParallelConfig), transmis à l’argument `parallel_config` de [from_pretrained()](/docs/diffusers/main/fr/api/models/overview#diffusers.ModelMixin.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](../training/distributed_inference#tensor-parallelism).
- l'encodeur de texte avec le [parallélisme de tenseurs](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) propre à Transformers, en passant un [DistributedConfig](https://huggingface.co/docs/transformers/main/en/expert_parallelism#transformers.DistributedConfig) et le même maillage.

Sur TPU, initialisez le groupe de processus avec `backend="tpu_dist"` et construisez le maillage avec `DeviceMesh("tpu", ...)`.

```python
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.

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

> [!WARNING]
> 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](https://github.com/google-pytorch/torch_tpu/) pour connaître la méthode actuelle de lancement des tâches multi-puces.