Google TPU

وحدات معالجة التينسور (TPUs) هي معجلات التعلم الآلي من Google، وهي متاحة على شكل Cloud TPUs (v6e, v7x). يعمل Diffusers عليها من خلال TorchTPU، وهو واجهة PyTorch الخاصة بـ TPUs، مع تغييرات برمجية طفيفة.

يتوفر ثلاثة أوضاع تنفيذ:

اتبع دليل تثبيت TorchTPU. بعد التثبيت، يُحمّل import torch المكتبة تلقائيًا ويسجّل جهاز "tpu"، لذا فإن pipe.to("tpu") هو التغيير الوحيد المطلوب. أضف import torch_tpu فقط إذا كنت قد عطّلت التحميل التلقائي للواجهة الخلفية باستخدام TORCH_DEVICE_BACKEND_AUTOLOAD=0.

الوضع الجشع

لا يتسع نموذج 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")

إذا كان النموذج كبيرًا جدًا ليُحمَّل على شريحة واحدة، أو كان لديك عدة شرائح وتريد تقليل زمن الاستجابة، فقسّم النماذج على الشرائح بدلاً من ذلك. راجع قسم التوازي على مستوى الموترات.

وضع التجميع

يسجّل TorchTPU "tpu" كاسم لخلفية torch.compile (تحت الغطاء TpuBackend)، لذا تُجمَّع المكونات كما تُجمَّع أي أهداف أخرى لـtorch.compile. يكون الاستدعاء الأول (إحماء) بطيئاً لأنه يقوم بالتجميع؛ أما الاستدعاءات اللاحقة ذات الأشكال نفسها فتعيد استخدام الرسم البياني المُجمَّع.

يتطلب TorchTPU أشكالاً ثابتة، لذا مرّر dynamic=False. يُعاد تجميع شكل جديد لـheight أو width مرة واحدة لكل شكل، بينما تُعاد استخدام الأشكال التي شوهدت مسبقاً. تغيير num_inference_steps لا يؤدي إلى إعادة التجميع.

عندما يتسع خط الأنابيب بالكامل على شريحة واحدة، انقله إلى TPU وادمج المحوّل الكامل. يتسع Stable Diffusion 3.5 Medium (حوالي 15GB في 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")

إذا لم يتسع خط المعالجة لشريحة واحدة:

• عند استخدام عدة شرائح، استخدم التوازي على مستوى المصفوفات بدلاً من ذلك. تبقى جميع العمليات على TPU، ويعمل pipe.transformer.compile(...) بنفس الطريقة على المحوّل المجزأ.
• مع enable_model_cpu_offload()، لا يمكن تتبّع خطافات النقل بواسطة torch.compile. قم بتجميع كتل التحوّل المتكررة فقط، باستخدام pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False).

توازي الموتر

قسّم النماذج التي تتجاوز سعة شريحة واحدة على عدة شرائح. يتجاوز كل من مُشفّر النص (~48GB) والمحوّل (~64GB) في 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)، وليس واجهة برمجة تطبيقات مستقرة، ومن المتوقع أن يتغير في إصدار مستقبلي من TorchTPU. تحقق من مستودع TorchTPU لمعرفة الطريقة الحالية لإطلاق مهام متعددة الرقائق.