# Google TPU

Tensor Processing Units (TPU) là bộ tăng tốc máy học của Google, có sẵn dưới dạng Cloud TPU (v6e, v7x). Diffusers chạy trên chúng thông qua [TorchTPU](https://github.com/google-pytorch/torch_tpu/), backend PyTorch cho TPU, với rất ít thay đổi mã.

Có ba chế độ thực thi:

| Chế độ | Hằng | Cách kích hoạt | Ghi chú |
|---|---|---|---|
| Strict eager (mặc định) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Các thao tác được thực hiện lần lượt một tại một thời điểm, bất đồng bộ |
| Eager được tối ưu | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | Các thao tác được trì hoãn và hợp nhất vào các đồ thị lớn hơn; nhanh hơn so với strict eager, thời gian biên dịch dài hơn |
| Biên dịch | — | `torch.compile(module, backend="tpu")` | Biên dịch AOT với `TpuBackend` |

Làm theo [hướng dẫn cài đặt TorchTPU](https://github.com/google-pytorch/torch_tpu/). Sau khi cài đặt, `import torch` sẽ tự động nạp và đăng ký thiết bị `"tpu"`, do đó `pipe.to("tpu")` là thay đổi duy nhất cần thực hiện. Chỉ thêm `import torch_tpu` nếu bạn đã tắt tự động nạp backend bằng `TORCH_DEVICE_BACKEND_AUTOLOAD=0`.

## Chế độ Eager

FLUX.1-schnell không vừa trên một chip v6e duy nhất cùng lúc, vì vậy hãy dùng [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/vi/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) để chuyển từng mô hình sang TPU chỉ trong khi nó chạy. Hàm này tự động phát hiện thiết bị `"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")
```

Nếu một model quá lớn cho một chip duy nhất, hoặc bạn có nhiều chip và muốn giảm độ trễ, hãy chia model thành các shard trên các chip thay vì chạy trên một chip. Xem phần [Tensor parallelism](#tensor-parallelism).

## Chế độ biên dịch

TorchTPU đăng ký `"tpu"` làm tên backend của `torch.compile` (bên trong là `TpuBackend`), do đó các thành phần được biên dịch giống như mọi mục tiêu `torch.compile` khác. Lệnh gọi đầu tiên (giai đoạn khởi động) sẽ chậm vì quá trình biên dịch; các lệnh gọi sau với cùng các hình dạng sẽ tái sử dụng đồ thị đã biên dịch.

> [!IMPORTANT]
> TorchTPU yêu cầu **hình dạng tĩnh**, vì vậy hãy truyền `dynamic=False`. Một `height` hoặc `width` mới sẽ được biên dịch lại cho hình dạng đó, chỉ một lần; các hình dạng đã gặp trước đó sẽ được tái sử dụng. Việc thay đổi `num_inference_steps` không gây ra biên dịch lại.

Khi toàn bộ pipeline vừa trên một chip, hãy di chuyển nó sang TPU và biên dịch toàn bộ transformer. Stable Diffusion 3.5 Medium (~15GB ở bf16) vừa trên một chip v6e duy nhất.

```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")
```

Nếu pipeline không vừa trên một chip:

- Với nhiều chip, hãy shard bằng [tensor parallelism](#tensor-parallelism) thay thế. Mọi thứ vẫn nằm trên TPU, và `pipe.transformer.compile(...)` hoạt động theo cách tương tự trên transformer đã shard.
- Với [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/vi/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload), các hook offload không thể được theo dõi bởi `torch.compile`. Thay vào đó, chỉ compile các khối lặp lại của transformer, bằng cách sử dụng `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`.

## Song song tensor

Tách các mô hình quá lớn cho một chip duy nhất thành nhiều chip. Bộ mã hóa văn bản (~48GB) và transformer (~64GB) của FLUX.2-dev đều vượt quá dung lượng của một chip, vì vậy ví dụ dưới đây tách cả hai:

- transformer với [TensorParallelConfig](/docs/diffusers/v0.43.0/vi/api/parallel#diffusers.TensorParallelConfig), được truyền vào đối số `parallel_config` của [from_pretrained()](/docs/diffusers/v0.43.0/vi/api/models/overview#diffusers.ModelMixin.from_pretrained). Mỗi rank chỉ đọc phần riêng của mình trong mỗi trọng số được chia nhỏ, nên mô hình đầy đủ không bao giờ được tạo ra. Để biết chi tiết chung về TP (`_tp_plan`, colwise/rowwise), hãy xem hướng dẫn [Tensor parallelism](./tensor_parallelism).
- bộ mã hóa văn bản với [song song tensor](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) của riêng Transformers, truyền vào một [DistributedConfig](https://huggingface.co/docs/transformers/v5.22.2/en/expert_parallelism#transformers.DistributedConfig) và cùng một lưới (mesh).

Trên TPU, khởi tạo nhóm tiến trình với `backend="tpu_dist"` và xây dựng lưới bằng `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")
```

Khởi chạy một tiến trình cho mỗi chip. Đặt `--nproc_per_node` để sử dụng toàn bộ số chip TPU trên máy chủ của bạn.

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

> [!WARNING]
> Lệnh khởi chạy ở trên sử dụng một module nội bộ của TorchTPU (`torch_tpu._internal`), không phải là API ổn định, và dự kiến sẽ thay đổi trong một bản phát hành TorchTPU tương lai. Hãy kiểm tra [kho TorchTPU](https://github.com/google-pytorch/torch_tpu/) để biết cách khởi chạy các tác vụ đa chip hiện tại.