Google TPU

Tensor Processing Units (TPU) — это ускорители машинного обучения от Google, доступные в виде Cloud TPU (v6e, v7x). Diffusers работает на них через TorchTPU, бэкенд PyTorch для TPU, с минимальными изменениями кода.

Доступны три режима выполнения:

Следуйте инструкции по установке TorchTPU. После установки import torch автоматически загружает его и регистрирует устройство "tpu", поэтому единственным необходимым изменением является pipe.to("tpu"). Добавляйте import torch_tpu только в том случае, если вы отключили автоматическую загрузку бэкенда с помощью TORCH_DEVICE_BACKEND_AUTOLOAD=0.

Режим 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")

Если модель слишком велика для одного чипа или у вас несколько чипов и вы хотите снизить задержку, распределите модели по чипам. См. раздел Tensor parallelism.

Скомпилированный режим

TorchTPU регистрирует "tpu" как имя бэкенда для torch.compile (внутри используется TpuBackend), поэтому компоненты компилируются так же, как и любые другие цели torch.compile. Первый вызов (прогрев) выполняется медленно, так как происходит компиляция; последующие вызовы с теми же формами повторно используют скомпилированный граф.

TorchTPU требует статических форм, поэтому передайте dynamic=False. Новая height или width приводит к повторной компиляции для этой формы один раз; уже встреченные формы переиспользуются. Изменение num_inference_steps не вызывает повторной компиляции.

Если весь пайплайн помещается на один чип, перенесите его на TPU и скомпилируйте полный трансформер. Stable Diffusion 3.5 Medium (~15 ГБ в bf16) помещается на один чип 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(...) работает так же и с шардированным трансформером.
• При использовании enable_model_cpu_offload() хуки offload не могут быть отслежены torch.compile. Компилируйте только повторяющиеся блоки трансформера, используя pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False).

Параллелизм тензоров

Размещайте модели, слишком большие для одного чипа, на нескольких. Текстовый энкодер (~48 ГБ) и трансформер (~64 ГБ) FLUX.2-dev по отдельности превышают объём одного чипа, поэтому в примере ниже оба компонента распределяются по нескольким чипам:

• трансформер с TensorParallelConfig, переданным в аргумент parallel_config метода from_pretrained(). Каждый ранг читает только свой срез каждого шардированного веса, поэтому полная модель никогда не материализуется. Общие сведения о TP (_tp_plan, colwise/rowwise) см. в руководстве Tensor parallelism.
• текстовый энкодер с собственным тензорным параллелизмом 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")

Запустите один процесс на чип. Установите --nproc_per_node в значение, равное общему числу чипов TPU на вашем хосте.

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.