Buckets:

hf-doc-build
/

doc

Files

xet

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

mishig

3 days ago

preview code

download

raw

8.53 kB

# Google TPU

Tensor Processing Units (TPU) は Google の機械学習アクセラレーターで、Cloud TPU (v6e、v7x) として利用できます。Diffusers は、TPU 向けの PyTorch バックエンドである [TorchTPU](https://github.com/google-pytorch/torch_tpu/) を通じて動作し、コードの修正は最小限で済みます。

3つの実行モードが利用可能です:

| モード | 定数 | 有効化する方法 | 注記 |
|---|---|---|---|
| 厳密な eager（デフォルト） | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | 1回ずつディスパッチされる非同期のオペレーション |
| 最適化された eager | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | 操作がデフェラされ、より大きなグラフに融合されます。厳密なイージアモードよりも高速ですが、コンパイル時間は長くなります。 |
| コンパイル | — | `torch.compile(module, backend="tpu")` | `TpuBackend` による AOT コンパイル |

[TorchTPU インストールガイド](https://github.com/google-pytorch/torch_tpu/)に従ってください。インストールが完了すると、`import torch` を実行するだけで自動的に読み込まれ、`"tpu"` デバイスが登録されるため、`pipe.to("tpu")` に変更するだけで対応できます。バックエンドの自動読み込みを `TORCH_DEVICE_BACKEND_AUTOLOAD=0` で無効化している場合のみ、`import torch_tpu` を追加してください。

## Eager モード

FLUX.1-schnell は単一の v6e チップに一度には収まらないため、[enable_model_cpu_offload()](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) を使用して、各モデルを実行している間のみ TPU へ移動させます。`"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")
```

モデルが1つのチップに収まらない場合、または複数のチップを所有していてレイテンシを下げたい場合は、代わりにモデルをチップ間でシャードしてください。[テンソル並列](#tensor-parallelism)セクションを参照してください。

## コンパイルモード

TorchTPU は `"tpu"` を `torch.compile` のバックエンド名として登録します（内部では `TpuBackend` として実装されています）。そのため、コンポーネントは他の `torch.compile` ターゲットと同様にコンパイルされます。最初の呼び出し（ウォームアップ）はコンパイルを行うため遅くなりますが、同じ形状での後続の呼び出しではコンパイル済みグラフが再利用されます。

> [!IMPORTANT]
> TorchTPU は**静的な形状**を必要とするため、`dynamic=False`を渡してください。新しい`height`または`width`が指定されると、その形状に対して一度だけ再コンパイルが行われます。以前に確認済みの形状は再利用されます。`num_inference_steps`を変更しても再コンパイルは行われません。

パイプライン全体が1つのチップに収まる場合は、TPUに移動してトランスフォーマー全体をコンパイルします。Stable Diffusion 3.5 Medium（bf16で約15GB）は、単一の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")
```

パイプラインが1つのチップに収まらない場合:

- 複数のチップを使用する場合は、代わりに[テンソル並列](#tensor-parallelism)でシャード化します。すべてがTPU上に保持され、`pipe.transformer.compile(...)`はシャード化されたtransformerでも同じように動作します。
- [enable_model_cpu_offload()](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) を使用すると、オフロードフックは `torch.compile` によってトレースできません。代わりに、`pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)` を使用して、transformer の繰り返しブロックのみをコンパイルしてください。

## テンソル並列

1つのチップに収まらないほど大きなモデルを複数のチップに分割します。FLUX.2-devのテキストエンコーダー（約48GB）とトランスフォーマー（約64GB）はそれぞれ1つのチップの容量を超えているため、以下の例では両方をシャードしています。

- [TensorParallelConfig](/docs/diffusers/main/ja/api/parallel#diffusers.TensorParallelConfig) を使用した transformer を、[from_pretrained()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.from_pretrained) の `parallel_config` 引数に渡します。各ランクは、シャード化された重みのうち自分のスライスのみを読み取るため、完全なモデルはメモリ上に展開されません。TP の一般的な詳細（`_tp_plan`、colwise/rowwise）については、[テンソル並列性](./tensor_parallelism)ガイドを参照してください。
- Transformers 独自の [テンソル並列](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) を使用してテキストエンコーダーを構成し、[DistributedConfig](https://huggingface.co/docs/transformers/main/en/expert_parallelism#transformers.DistributedConfig) と同じメッシュを渡します。

TPU では、`backend="tpu_dist"` を指定してプロセスグループを初期化し、`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")
```

チップごとに1つのプロセスを起動します。ホスト上のすべてのTPUチップを使用するには、`--nproc_per_node`を設定してください。

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

> [!WARNING]
> 上記の起動コマンドは、安定した API ではなく TorchTPU の内部モジュール（`torch_tpu._internal`）を使用しており、将来の TorchTPU リリースで変更される予定です。マルチチップジョブの現在の起動方法については、[TorchTPU リポジトリ](https://github.com/google-pytorch/torch_tpu/)を確認してください。

Xet Storage Details

Size:: 8.53 kB
Xet hash:: e6fd6e5a3b67c636c24d7d9b7611b54b3cd9023b2e7d840acb0f236015661bb2

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

	# Google TPU

	Tensor Processing Units (TPU) は Google の機械学習アクセラレーターで、Cloud TPU (v6e、v7x) として利用できます。Diffusers は、TPU 向けの PyTorch バックエンドである [TorchTPU](https://github.com/google-pytorch/torch_tpu/) を通じて動作し、コードの修正は最小限で済みます。

	3つの実行モードが利用可能です:

	\| モード \| 定数 \| 有効化する方法 \| 注記 \|
	\|---\|---\|---\|---\|
	\| 厳密な eager（デフォルト） \| `EagerMode.DEFER_NEVER` \| `pipe.to("tpu")` \| 1回ずつディスパッチされる非同期のオペレーション \|
	\| 最適化された eager \| `EagerMode.DEFER_AND_FUSE` \| `TORCH_TPU_DEFER_AND_FUSE=1` \| 操作がデフェラされ、より大きなグラフに融合されます。厳密なイージアモードよりも高速ですが、コンパイル時間は長くなります。 \|
	\| コンパイル \| — \| `torch.compile(module, backend="tpu")` \| `TpuBackend` による AOT コンパイル \|

	[TorchTPU インストールガイド](https://github.com/google-pytorch/torch_tpu/)に従ってください。インストールが完了すると、`import torch` を実行するだけで自動的に読み込まれ、`"tpu"` デバイスが登録されるため、`pipe.to("tpu")` に変更するだけで対応できます。バックエンドの自動読み込みを `TORCH_DEVICE_BACKEND_AUTOLOAD=0` で無効化している場合のみ、`import torch_tpu` を追加してください。

	## Eager モード

	FLUX.1-schnell は単一の v6e チップに一度には収まらないため、[enable_model_cpu_offload()](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) を使用して、各モデルを実行している間のみ TPU へ移動させます。`"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")
	```

	モデルが1つのチップに収まらない場合、または複数のチップを所有していてレイテンシを下げたい場合は、代わりにモデルをチップ間でシャードしてください。[テンソル並列](#tensor-parallelism)セクションを参照してください。

	## コンパイルモード

	TorchTPU は `"tpu"` を `torch.compile` のバックエンド名として登録します（内部では `TpuBackend` として実装されています）。そのため、コンポーネントは他の `torch.compile` ターゲットと同様にコンパイルされます。最初の呼び出し（ウォームアップ）はコンパイルを行うため遅くなりますが、同じ形状での後続の呼び出しではコンパイル済みグラフが再利用されます。

	> [!IMPORTANT]
	> TorchTPU は静的な形状を必要とするため、`dynamic=False`を渡してください。新しい`height`または`width`が指定されると、その形状に対して一度だけ再コンパイルが行われます。以前に確認済みの形状は再利用されます。`num_inference_steps`を変更しても再コンパイルは行われません。

	パイプライン全体が1つのチップに収まる場合は、TPUに移動してトランスフォーマー全体をコンパイルします。Stable Diffusion 3.5 Medium（bf16で約15GB）は、単一の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")
	```

	パイプラインが1つのチップに収まらない場合:

	- 複数のチップを使用する場合は、代わりに[テンソル並列](#tensor-parallelism)でシャード化します。すべてがTPU上に保持され、`pipe.transformer.compile(...)`はシャード化されたtransformerでも同じように動作します。
	- [enable_model_cpu_offload()](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) を使用すると、オフロードフックは `torch.compile` によってトレースできません。代わりに、`pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)` を使用して、transformer の繰り返しブロックのみをコンパイルしてください。

	## テンソル並列

	1つのチップに収まらないほど大きなモデルを複数のチップに分割します。FLUX.2-devのテキストエンコーダー（約48GB）とトランスフォーマー（約64GB）はそれぞれ1つのチップの容量を超えているため、以下の例では両方をシャードしています。

	- [TensorParallelConfig](/docs/diffusers/main/ja/api/parallel#diffusers.TensorParallelConfig) を使用した transformer を、[from_pretrained()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.from_pretrained) の `parallel_config` 引数に渡します。各ランクは、シャード化された重みのうち自分のスライスのみを読み取るため、完全なモデルはメモリ上に展開されません。TP の一般的な詳細（`_tp_plan`、colwise/rowwise）については、[テンソル並列性](./tensor_parallelism)ガイドを参照してください。
	- Transformers 独自の [テンソル並列](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) を使用してテキストエンコーダーを構成し、[DistributedConfig](https://huggingface.co/docs/transformers/main/en/expert_parallelism#transformers.DistributedConfig) と同じメッシュを渡します。

	TPU では、`backend="tpu_dist"` を指定してプロセスグループを初期化し、`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")
	```

	チップごとに1つのプロセスを起動します。ホスト上のすべてのTPUチップを使用するには、`--nproc_per_node`を設定してください。

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

	> [!WARNING]
	> 上記の起動コマンドは、安定した API ではなく TorchTPU の内部モジュール（`torch_tpu._internal`）を使用しており、将来の TorchTPU リリースで変更される予定です。マルチチップジョブの現在の起動方法については、[TorchTPU リポジトリ](https://github.com/google-pytorch/torch_tpu/)を確認してください。