diffusers/main/ja/optimization/tpu.md

Google TPU

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

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

Eager モード

FLUX.1-schnell は単一の v6e チップに一度には収まらないため、enable_model_cpu_offload() を使用して、各モデルを実行している間のみ TPU へ移動させます。"tpu" デバイスは自動的に検出されます。

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つのチップに収まらない場合、または複数のチップを所有していてレイテンシを下げたい場合は、代わりにモデルをチップ間でシャードしてください。テンソル並列セクションを参照してください。

コンパイルモード

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

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

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

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

• 複数のチップを使用する場合は、代わりにテンソル並列でシャード化します。すべてがTPU上に保持され、pipe.transformer.compile(...)はシャード化されたtransformerでも同じように動作します。
• 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 を使用した transformer を、from_pretrained() の parallel_config 引数に渡します。各ランクは、シャード化された重みのうち自分のスライスのみを読み取るため、完全なモデルはメモリ上に展開されません。TP の一般的な詳細（_tp_plan、colwise/rowwise）については、テンソル並列性ガイドを参照してください。
• Transformers 独自の テンソル並列 を使用してテキストエンコーダーを構成し、DistributedConfig と同じメッシュを渡します。

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

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

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

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