Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /v0.43.0 /pt /optimization /tpu.md

mishig

2 days ago

preview code

download

raw

7.17 kB

# Google TPU

Tensor Processing Units (TPUs) são os aceleradores de machine learning da Google, disponíveis como Cloud TPUs (v6e, v7x). O Diffusers roda neles por meio do [TorchTPU](https://github.com/google-pytorch/torch_tpu/), o backend do PyTorch para TPUs, com mudanças mínimas no código.

Três modos de execução estão disponíveis:

| Modo | Constante | Como ativar | Observações |
|---|---|---|---|
| Strict eager (padrão) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Operações despachadas uma de cada vez, assíncronas |
| Eager otimizado | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | Operações adiadas e fundidas em grafos maiores; mais rápido que o modo eager estrito, com tempo de compilação mais longo |
| Compilar | — | `torch.compile(module, backend="tpu")` | Compilação AOT com `TpuBackend` |

Siga o [guia de instalação do TorchTPU](https://github.com/google-pytorch/torch_tpu/). Uma vez instalado, `import torch` o carrega automaticamente e registra o dispositivo `"tpu"`, de modo que `pipe.to("tpu")` é a única alteração necessária. Adicione `import torch_tpu` apenas se você desativou o carregamento automático do backend com `TORCH_DEVICE_BACKEND_AUTOLOAD=0`.

## Modo Eager

FLUX.1-schnell não cabe em um único chip v6e de uma vez, então use [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/pt/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) para mover cada modelo para o TPU apenas enquanto ele está em execução. Ele detecta automaticamente o 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")
```

Se um modelo for grande demais para um único chip, ou se você tiver vários chips e quiser menor latência, distribua os modelos entre os chips. Consulte a seção [Paralelismo de tensor](#tensor-parallelism).

## Modo compilado

TorchTPU registra `"tpu"` como um nome de backend para `torch.compile` (`TpuBackend` por baixo dos panos), de modo que os componentes são compilados como qualquer outro alvo de `torch.compile`. A primeira chamada (aquecimento) é lenta porque realiza a compilação; chamadas subsequentes com as mesmas formas reutilizam o gráfico compilado.

> [!IMPORTANT]
> TorchTPU requer **formas estáticas**, então passe `dynamic=False`. Um novo `height` ou `width` compila novamente para aquela forma, uma vez; formas já vistas são reutilizadas. Alterar `num_inference_steps` não recompila.

Quando o pipeline inteiro cabe em um único chip, mova-o para o TPU e compile o transformer completo. O Stable Diffusion 3.5 Medium (~15GB em bf16) cabe em um único 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")
```

Se o pipeline não couber em um único chip:

- Com vários chips, use [tensor parallelism](#tensor-parallelism) para fazer o sharding. Tudo permanece no TPU e `pipe.transformer.compile(...)` funciona da mesma forma no transformer shardado.
- Com [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/pt/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload), os hooks de offload não podem ser rastreados pelo `torch.compile`. Compile apenas os blocos repetidos do transformer, com `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`.

## Paralelismo de tensor

Divida modelos grandes demais para um único chip entre vários. O codificador de texto (~48GB) e o transformer (~64GB) do FLUX.2-dev excedem a capacidade de um único chip, portanto, o exemplo abaixo divide ambos:

- o transformer com [TensorParallelConfig](/docs/diffusers/v0.43.0/pt/api/parallel#diffusers.TensorParallelConfig), passado para o argumento `parallel_config` de [from_pretrained()](/docs/diffusers/v0.43.0/pt/api/models/overview#diffusers.ModelMixin.from_pretrained). Cada rank lê apenas a própria fatia de cada peso particionado, de modo que o modelo completo nunca é materializado. Para detalhes gerais de TP (`_tp_plan`, colwise/rowwise), consulte o guia [Tensor parallelism](./tensor_parallelism).
- o codificador de texto com o próprio [paralelismo de tensor](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) dos Transformers, passando um [DistributedConfig](https://huggingface.co/docs/transformers/v5.22.2/en/expert_parallelism#transformers.DistributedConfig) e a mesma malha.

No TPU, inicialize o grupo de processos com `backend="tpu_dist"` e construa a malha com `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 um processo por chip. Defina `--nproc_per_node` para usar o número total de chips TPU no seu host.

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

> [!WARNING]
> O comando de inicialização acima usa um módulo interno do TorchTPU (`torch_tpu._internal`), não uma API estável, e é esperado que mude em uma futura liberação do TorchTPU. Consulte o [repositório do TorchTPU](https://github.com/google-pytorch/torch_tpu/) para a maneira atual de iniciar trabalhos multi-chip.

Xet Storage Details

Size:: 7.17 kB
Xet hash:: 6adfae2c0f2bc789a5c9144ddf5044a49f29d99f22eb058c41735e1e39e2fdfc

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

	# Google TPU

	Tensor Processing Units (TPUs) são os aceleradores de machine learning da Google, disponíveis como Cloud TPUs (v6e, v7x). O Diffusers roda neles por meio do [TorchTPU](https://github.com/google-pytorch/torch_tpu/), o backend do PyTorch para TPUs, com mudanças mínimas no código.

	Três modos de execução estão disponíveis:

	\| Modo \| Constante \| Como ativar \| Observações \|
	\|---\|---\|---\|---\|
	\| Strict eager (padrão) \| `EagerMode.DEFER_NEVER` \| `pipe.to("tpu")` \| Operações despachadas uma de cada vez, assíncronas \|
	\| Eager otimizado \| `EagerMode.DEFER_AND_FUSE` \| `TORCH_TPU_DEFER_AND_FUSE=1` \| Operações adiadas e fundidas em grafos maiores; mais rápido que o modo eager estrito, com tempo de compilação mais longo \|
	\| Compilar \| — \| `torch.compile(module, backend="tpu")` \| Compilação AOT com `TpuBackend` \|

	Siga o [guia de instalação do TorchTPU](https://github.com/google-pytorch/torch_tpu/). Uma vez instalado, `import torch` o carrega automaticamente e registra o dispositivo `"tpu"`, de modo que `pipe.to("tpu")` é a única alteração necessária. Adicione `import torch_tpu` apenas se você desativou o carregamento automático do backend com `TORCH_DEVICE_BACKEND_AUTOLOAD=0`.

	## Modo Eager

	FLUX.1-schnell não cabe em um único chip v6e de uma vez, então use [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/pt/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) para mover cada modelo para o TPU apenas enquanto ele está em execução. Ele detecta automaticamente o 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")
	```

	Se um modelo for grande demais para um único chip, ou se você tiver vários chips e quiser menor latência, distribua os modelos entre os chips. Consulte a seção [Paralelismo de tensor](#tensor-parallelism).

	## Modo compilado

	TorchTPU registra `"tpu"` como um nome de backend para `torch.compile` (`TpuBackend` por baixo dos panos), de modo que os componentes são compilados como qualquer outro alvo de `torch.compile`. A primeira chamada (aquecimento) é lenta porque realiza a compilação; chamadas subsequentes com as mesmas formas reutilizam o gráfico compilado.

	> [!IMPORTANT]
	> TorchTPU requer formas estáticas, então passe `dynamic=False`. Um novo `height` ou `width` compila novamente para aquela forma, uma vez; formas já vistas são reutilizadas. Alterar `num_inference_steps` não recompila.

	Quando o pipeline inteiro cabe em um único chip, mova-o para o TPU e compile o transformer completo. O Stable Diffusion 3.5 Medium (~15GB em bf16) cabe em um único 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")
	```

	Se o pipeline não couber em um único chip:

	- Com vários chips, use [tensor parallelism](#tensor-parallelism) para fazer o sharding. Tudo permanece no TPU e `pipe.transformer.compile(...)` funciona da mesma forma no transformer shardado.
	- Com [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/pt/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload), os hooks de offload não podem ser rastreados pelo `torch.compile`. Compile apenas os blocos repetidos do transformer, com `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`.

	## Paralelismo de tensor

	Divida modelos grandes demais para um único chip entre vários. O codificador de texto (~48GB) e o transformer (~64GB) do FLUX.2-dev excedem a capacidade de um único chip, portanto, o exemplo abaixo divide ambos:

	- o transformer com [TensorParallelConfig](/docs/diffusers/v0.43.0/pt/api/parallel#diffusers.TensorParallelConfig), passado para o argumento `parallel_config` de [from_pretrained()](/docs/diffusers/v0.43.0/pt/api/models/overview#diffusers.ModelMixin.from_pretrained). Cada rank lê apenas a própria fatia de cada peso particionado, de modo que o modelo completo nunca é materializado. Para detalhes gerais de TP (`_tp_plan`, colwise/rowwise), consulte o guia [Tensor parallelism](./tensor_parallelism).
	- o codificador de texto com o próprio [paralelismo de tensor](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) dos Transformers, passando um [DistributedConfig](https://huggingface.co/docs/transformers/v5.22.2/en/expert_parallelism#transformers.DistributedConfig) e a mesma malha.

	No TPU, inicialize o grupo de processos com `backend="tpu_dist"` e construa a malha com `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 um processo por chip. Defina `--nproc_per_node` para usar o número total de chips TPU no seu host.

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

	> [!WARNING]
	> O comando de inicialização acima usa um módulo interno do TorchTPU (`torch_tpu._internal`), não uma API estável, e é esperado que mude em uma futura liberação do TorchTPU. Consulte o [repositório do TorchTPU](https://github.com/google-pytorch/torch_tpu/) para a maneira atual de iniciar trabalhos multi-chip.