Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /main /fr /optimization /tpu.md

mishig

3 days ago

preview code

download

raw

7.52 kB

# 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](./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.

Xet Storage Details

Size:: 7.52 kB
Xet hash:: aa1dd24375ac24ee6e7e169e5107e3deea6c3822a250232c2ed3947418af0aa7

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	# 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](./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.