Buckets:

hf-doc-build
/

doc

Files

xet

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

mishig

3 days ago

preview code

download

raw

7.17 kB

# 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](https://github.com/google-pytorch/torch_tpu/), 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](https://github.com/google-pytorch/torch_tpu/). 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()](/docs/diffusers/main/id/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) untuk memindahkan setiap model ke TPU hanya saat model tersebut berjalan. Fungsi ini mendeteksi perangkat `"tpu"` secara otomatis.

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

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

> [!IMPORTANT]
> 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.

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

Jika pipeline tidak muat di satu chip:

- Dengan beberapa chip, shard-kan dengan [tensor parallelism](#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()](/docs/diffusers/main/id/api/pipelines/overview#diffusers.DiffusionPipeline.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](/docs/diffusers/main/id/api/parallel#diffusers.TensorParallelConfig), yang diteruskan ke argumen `parallel_config` dari [from_pretrained()](/docs/diffusers/main/id/api/models/overview#diffusers.ModelMixin.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](./tensor_parallelism).
- text encoder dengan [parallelisme tensor](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) milik Transformers, dengan meneruskan [DistributedConfig](https://huggingface.co/docs/transformers/main/en/expert_parallelism#transformers.DistributedConfig) dan mesh yang sama.

Pada TPU, inisialisasi grup proses dengan `backend="tpu_dist"` dan bangun mesh dengan `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")
```

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

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

> [!WARNING]
> 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](https://github.com/google-pytorch/torch_tpu/) untuk cara terkini meluncurkan pekerjaan multi-chip.

Xet Storage Details

Size:: 7.17 kB
Xet hash:: f8776be63326453e1d56fa24c3031afc4195acac99c7c71e0f2827c54abefd67

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

	# 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](https://github.com/google-pytorch/torch_tpu/), 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](https://github.com/google-pytorch/torch_tpu/). 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()](/docs/diffusers/main/id/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) untuk memindahkan setiap model ke TPU hanya saat model tersebut berjalan. Fungsi ini mendeteksi perangkat `"tpu"` secara otomatis.

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

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

	> [!IMPORTANT]
	> 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.

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

	Jika pipeline tidak muat di satu chip:

	- Dengan beberapa chip, shard-kan dengan [tensor parallelism](#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()](/docs/diffusers/main/id/api/pipelines/overview#diffusers.DiffusionPipeline.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](/docs/diffusers/main/id/api/parallel#diffusers.TensorParallelConfig), yang diteruskan ke argumen `parallel_config` dari [from_pretrained()](/docs/diffusers/main/id/api/models/overview#diffusers.ModelMixin.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](./tensor_parallelism).
	- text encoder dengan [parallelisme tensor](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) milik Transformers, dengan meneruskan [DistributedConfig](https://huggingface.co/docs/transformers/main/en/expert_parallelism#transformers.DistributedConfig) dan mesh yang sama.

	Pada TPU, inisialisasi grup proses dengan `backend="tpu_dist"` dan bangun mesh dengan `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")
	```

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

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

	> [!WARNING]
	> 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](https://github.com/google-pytorch/torch_tpu/) untuk cara terkini meluncurkan pekerjaan multi-chip.