diffusers/main/ko/optimization/tpu.md

Google TPU

Tensor Processing Units(TPU)는 Google의 머신러닝 가속기로, Cloud TPU(v6e, v7x)로 제공됩니다. Diffusers는 TPU용 PyTorch 백엔드인 TorchTPU를 통해 TPU에서 최소한의 코드 변경으로 실행됩니다.

세 가지 실행 모드가 사용할 수 있습니다:

| 모드 | 상수 | 활성화 방법 | The input is a single word "Notes".

참고 | |---|---|---|---| | 엄격한 eager (기본값) | EagerMode.DEFER_NEVER | pipe.to("tpu") | Operations dispatched one at a time, asynchronous | | 최적화된 eager | EagerMode.DEFER_AND_FUSE | TORCH_TPU_DEFER_AND_FUSE=1 | 작업이 지연되고 더 큰 그래프로 융합됩니다. 엄격한 eager보다 빠르지만 컴파일이 더 오래 걸립니다 | | 컴파일 | — | 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 칩에 한 번에 모두 들어가지 않으므로, 각 모델이 실행되는 동안에만 TPU로 이동하도록 enable_model_cpu_offload()를 사용하세요. "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")

모델이 단일 칩에 너무 크거나, 여러 칩을 보유하고 있어 낮은 지연 시간을 원한다면, 대신 모델을 칩에 걸쳐 샤딩하세요. Tensor parallelism 섹션을 참조하세요.

컴파일 모드

TorchTPU는 "tpu"를 torch.compile 백엔드 이름으로 등록합니다(내부적으로는 TpuBackend를 사용하므로), 따라서 각 구성 요소는 다른 torch.compile 대상과 동일한 방식으로 컴파일됩니다. 첫 번째 호출(워밍업)은 컴파일을 수행하기 때문에 느리지만, 동일한 shape을 가진 후속 호출은 컴파일된 그래프를 재사용합니다.

TorchTPU는 정적 shape를 요구하므로 dynamic=False를 전달해야 합니다. 새로운 height 또는 width는 해당 shape에 대해 한 번만 다시 컴파일되며, 이미 확인된 shape는 재사용됩니다. num_inference_steps를 변경해도 다시 컴파일되지 않습니다.

파이프라인 전체가 하나의 칩에 들어갈 수 있는 경우, 이를 TPU로 이동하고 전체 transformer를 컴파일합니다. 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")

파이프라인이 하나의 칩에 맞지 않는 경우:

• 여러 개의 칩을 사용하는 경우, 대신 tensor parallelism으로 샤딩하세요. 모든 것이 TPU에 유지되며, pipe.transformer.compile(...)은 샤딩된 transformer에서도 동일한 방식으로 작동합니다.
• enable_model_cpu_offload()를 사용하면 오프로드 훅을 torch.compile로 추적할 수 없습니다. 대신 pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)를 사용하여 transformer의 반복 블록만 컴파일하세요.

Tensor 병렬성

하나의 칩에 너무 큰 모델을 여러 칩에 걸쳐 샤딩합니다. FLUX.2-dev의 텍스트 인코더(48GB)와 트랜스포머(64GB)는 각각 단일 칩을 초과하므로, 아래 예시에서는 둘 다 샤딩합니다:

• parallel_config 인수로 전달되는 from_pretrained()에 TensorParallelConfig를 갖춘 transformer입니다. 각 rank는 샤딩된 가중치의 자기 자신의 슬라이스만 읽으므로 전체 모델은 절대 메모리에 로드되지 않습니다. 일반적인 TP 세부 사항(_tp_plan, colwise/rowwise)은 Tensor parallelism 가이드를 참조하세요.
• Transformers 자체의 tensor parallelism를 사용하여 텍스트 인코더를 구성하고, DistributedConfig와 동일한 mesh를 전달합니다.

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

칩당 하나의 프로세스를 시작합니다. 호스트의 모든 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

위 시작 명령은 TorchTPU 내부 모듈(torch_tpu._internal)을 사용하며, 이는 안정화된 API가 아니므로 향후 TorchTPU 릴리스에서 변경될 수 있습니다. 멀티 칩 작업을 시작하는 최신 방법은 TorchTPU 저장소에서 확인하세요.