diffusers/main/id/optimization/tpu.md

Google TPU

Tensor Processing Units (TPU) adalah akselerator machine learning dari Google, yang tersedia sebagai Cloud TPU (v6e, v7x). Diffusers berjalan di atasnya melalui TorchTPU, backend PyTorch untuk TPU, dengan perubahan kode yang minimal.

Tiga mode eksekusi tersedia:

Mode	Konstan	Cara mengaktifkan	Catatan
Strict eager (default)	EagerMode.DEFER_NEVER	pipe.to("tpu")	Operasi yang dikirim satu per satu, asinkron
Eager yang dioptimalkan	EagerMode.DEFER_AND_FUSE	TORCH_TPU_DEFER_AND_FUSE=1	Operasi ditunda dan dilebur ke dalam grafik yang lebih besar; lebih cepat daripada strict eager, waktu kompilasi lebih lama
Compile	—	torch.compile(module, backend="tpu")	Kompilasi AOT dengan TpuBackend

Ikuti panduan instalasi TorchTPU. Setelah terinstal, import torch akan memuatnya secara otomatis dan mendaftarkan perangkat "tpu", sehingga pipe.to("tpu") adalah satu-satunya perubahan yang diperlukan. Tambahkan import torch_tpu hanya jika Anda menonaktifkan pemuatan otomatis backend dengan TORCH_DEVICE_BACKEND_AUTOLOAD=0.

Mode Eager

FLUX.1-schnell tidak muat di satu chip v6e sekaligus, jadi gunakan enable_model_cpu_offload() untuk memindahkan setiap model ke TPU hanya saat model tersebut berjalan. Fungsi ini mendeteksi perangkat "tpu" secara otomatis.

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

Jika model terlalu besar untuk satu chip, atau Anda memiliki beberapa chip dan ingin latensi yang lebih rendah, pecah model di antara chip tersebut. Lihat bagian Tensor parallelism.

Mode terkompilasi

TorchTPU mendaftarkan "tpu" sebagai nama backend torch.compile (TpuBackend di baliknya), sehingga komponen dikompilasi seperti target torch.compile lainnya. Panggilan pertama (warmup) lambat karena proses kompilasi; panggilan berikutnya dengan bentuk (shapes) yang sama akan menggunakan ulang graf yang telah dikompilasi.

TorchTPU memerlukan bentuk statis, jadi berikan dynamic=False. Nilai height atau width baru akan dikompilasi ulang untuk bentuk tersebut, sekali; bentuk yang sudah pernah dilihat akan digunakan kembali. Mengubah num_inference_steps tidak memicu kompilasi ulang.

Ketika seluruh pipeline muat di satu chip, pindahkan ke TPU dan kompilasi transformer lengkap. Stable Diffusion 3.5 Medium (~15GB dalam bf16) muat di satu chip v6e.

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

Jika pipeline tidak muat di satu chip:

• Dengan beberapa chip, shard-kan dengan tensor parallelism sebagai gantinya. Semua tetap berada di TPU, dan pipe.transformer.compile(...) bekerja dengan cara yang sama pada transformer yang di-shard.
• Dengan enable_model_cpu_offload(), hook offload tidak dapat dilacak oleh torch.compile. Kompilasi hanya blok berulang dari transformer, dengan pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False).

Paralelisme tensor

Shard model yang terlalu besar untuk satu chip ke beberapa chip. Text encoder FLUX.2-dev (48GB) dan transformer (64GB) masing-masing melebihi kapasitas satu chip, sehingga contoh di bawah ini melakukan sharding pada keduanya:

• transformer dengan TensorParallelConfig, yang diteruskan ke argumen parallel_config dari from_pretrained(). Setiap rank hanya membaca slice miliknya sendiri dari setiap bobot yang di-shard, sehingga model lengkap tidak pernah dimaterialisasi. Untuk detail TP umum (_tp_plan, colwise/rowwise), lihat panduan Tensor parallelism.
• text encoder dengan parallelisme tensor milik Transformers, dengan meneruskan DistributedConfig dan mesh yang sama.

Pada TPU, inisialisasi grup proses dengan backend="tpu_dist" dan bangun mesh dengan DeviceMesh("tpu", ...).

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

Luncurkan satu proses per chip. Atur --nproc_per_node untuk menggunakan seluruh jumlah chip TPU di host Anda.

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

Perintah peluncuran di atas menggunakan modul internal TorchTPU (torch_tpu._internal), bukan API yang stabil, dan diperkirakan akan berubah pada rilis TorchTPU di masa depan. Periksa repo TorchTPU untuk cara terkini meluncurkan pekerjaan multi-chip.