# 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](https://github.com/google-pytorch/torch_tpu/), el backend de PyTorch para TPU, con cambios mínimos en el código.

Hay disponibles tres modos de ejecución:

| Modo | Constante | Cómo activar | Notas |
|---|---|---|---|
| Eager estricto (predeterminado) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Operaciones despachadas una a la vez, asíncronas |
| Eager optimizado | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | Operaciones diferidas y fusionadas en grafos más grandes; más rápido que el modo estricto eager, con una compilación más larga |
| Compilar | — | `torch.compile(module, backend="tpu")` | Compilación AOT con `TpuBackend` |

Siga la [guía de instalación de TorchTPU](https://github.com/google-pytorch/torch_tpu/). 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()](/docs/diffusers/v0.43.0/es/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) para mover cada modelo al TPU solo mientras se ejecuta. Detecta automáticamente el dispositivo `"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 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](#tensor-parallelism).

## 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.

> [!IMPORTANT]
> 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 transformer completo. Stable Diffusion 3.5 Medium (~15GB en bf16) cabe en un solo chip 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 el pipeline no cabe en un solo chip:

- Con varios chips, utilice [tensor parallelism](#tensor-parallelism) para fragmentarlo. Todo permanece en la TPU y `pipe.transformer.compile(...)` funciona de la misma manera en el transformer fragmentado.
- Con [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/es/api/pipelines/overview#diffusers.DiffusionPipeline.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 en varios chips los modelos demasiado grandes para un solo chip. El codificador de texto (~48GB) y el transformador (~64GB) de FLUX.2-dev superan la capacidad de un chip individual, por lo que el ejemplo siguiente divide ambos:

- el transformer con [TensorParallelConfig](/docs/diffusers/v0.43.0/es/api/parallel#diffusers.TensorParallelConfig), pasado al argumento `parallel_config` de [from_pretrained()](/docs/diffusers/v0.43.0/es/api/models/overview#diffusers.ModelMixin.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](./tensor_parallelism).
- el codificador de texto con el propio [paralelismo de tensores](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) de Transformers, pasando un [DistributedConfig](https://huggingface.co/docs/transformers/v5.22.2/en/expert_parallelism#transformers.DistributedConfig) y la misma malla.

En TPU, inicializa el grupo de procesos con `backend="tpu_dist"` y construye la malla con `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")
```

Inicie un proceso por chip. Establezca `--nproc_per_node` para usar el número total de chips TPU de su host.

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

> [!WARNING]
> 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](https://github.com/google-pytorch/torch_tpu/) para conocer la forma actual de lanzar trabajos de varios chips.