Google TPU

Tensor Processing Units (TPU'lar), Google'ın makine öğrenmesi hızlandırıcılarıdır ve Cloud TPU (v6e, v7x) olarak kullanılabilir. Diffusers, TPU'lar için PyTorch arka ucu olan TorchTPU aracılığıyla bunlarda çalışır ve minimal kod değişikliği gerektirir.

Üç çalışma modu mevcuttur:

TorchTPU kurulum kılavuzunu izleyin. Kurulum tamamlandıktan sonra import torch otomatik olarak yüklenir ve "tpu" cihazını kaydeder; bu nedenle tek gereken değişiklik pipe.to("tpu") komutunu eklemektir. TORCH_DEVICE_BACKEND_AUTOLOAD=0 ile arka uç otomatik yüklemeyi devre dışı bıraktıysanız, yalnızca import torch_tpu satırını ekleyin.

Eager mod

FLUX.1-schnell tek bir v6e çipine bir arada sığmadığı için, her modelin yalnızca çalıştığı sırada TPU'ya taşınmasını sağlamak üzere enable_model_cpu_offload() kullanın. "tpu" cihazını otomatik olarak algılar.

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

Bir model tek bir çip için çok büyükse veya birden fazla çipiniz varsa ve daha düşük gecikme süresi istiyorsanız, modelleri çipler arasında bölüştürün. Tensor parallelism bölümüne bakın.

Derlenmiş mod

TorchTPU, "tpu" adını torch.compile arka uç adı olarak kaydeder (altta TpuBackend kullanılır), bu nedenle bileşenler herhangi bir torch.compile hedefi gibi derlenir. İlk çağrı (ısınma) derleme nedeniyle yavaştır; aynı şekillerle yapılan sonraki çağrılar derlenmiş grafiği yeniden kullanır.

TorchTPU statik şekiller gerektirir, bu nedenle dynamic=False parametresini iletin. Yeni bir height veya width değeri, o şekil için bir kez yeniden derlenir; daha önce görülmüş şekiller yeniden kullanılır. num_inference_steps değerini değiştirmek yeniden derlemeye neden olmaz.

Tüm pipeline tek bir çipe sığıyorsa, onu TPU'ya taşıyın ve tam dönüştürücüyü derleyin. Stable Diffusion 3.5 Orta (~15GB bf16'da) tek bir v6e çipine sığar.

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

Eğer pipeline tek bir çipe sığmıyorsa:

• Birden fazla çip kullanıldığında, bunun yerine tensor parallelism ile parçalayın. Her şey TPU üzerinde kalır ve pipe.transformer.compile(...), parçalanmış transformer üzerinde aynı şekilde çalışır.
• enable_model_cpu_offload() ile offload kancaları torch.compile tarafından izlenemez. Bunun yerine yalnızca transformer’ın tekrarlanan bloklarını derleyin: pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False).

Tensor paralelliği

Tek bir çipe sığmayacak kadar büyük modelleri birden fazla çipe dağıtın. FLUX.2-dev'in metin kodlayıcısı (~48GB) ve dönüştürücüsü (~64GB) tek bir çipi aşar, bu nedenle aşağıdaki örnek her ikisini de parçalar:

• TensorParallelConfig ile transformer, from_pretrained() fonksiyonunun parallel_config argümanına geçirilir. Her rank, shard edilmiş her ağırlığın yalnızca kendi dilimini okuduğu için tam model asla maddi olarak oluşturulmaz. Genel TP ayrıntıları için (_tp_plan, colwise/rowwise) Tensor parallelism kılavuzuna bakın.
• Transformers'ın kendi tensor parallelism'i ile metin kodlayıcıyı, bir DistributedConfig ve aynı mesh'i geçirerek.

TPU'da, süreç grubunu backend="tpu_dist" ile başlatın ve mesh'i DeviceMesh("tpu", ...) ile oluşturun.

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

Her çip için bir süreç başlatın. --nproc_per_node parametresini, hostunuzdaki tüm TPU çiplerini kullanmak üzere ayarlayın.

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

Yukarıdaki başlatma komutu, TorchTPU iç modülü olan torch_tpu._internal'i kullanır; bu bir kararlı API değildir ve gelecekteki bir TorchTPU sürümünde değişmesi beklenmektedir. Çok çipli işleri başlatmanın güncel yöntemi için TorchTPU deposunu kontrol edin.