# Google TPU

وحدات معالجة التينسور (TPUs) هي معجلات التعلم الآلي من Google، وهي متاحة على شكل Cloud TPUs (v6e, v7x). يعمل Diffusers عليها من خلال [TorchTPU](https://github.com/google-pytorch/torch_tpu/)، وهو واجهة PyTorch الخاصة بـ TPUs، مع تغييرات برمجية طفيفة.

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

| الوضع | ثابت | كيفية التفعيل | ملاحظات |
|---|---|---|---|
| الحماس الصارم (الافتراضي) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | عمليات تُنفَّذ واحدة تلو الأخرى، غير متزامنة |
| الحماسي المُحسَّن | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | عمليات مؤجلة ومدمجة في رسوم بيانية أكبر؛ أسرع من الوضع العاجل الصارم، مع وقت تجميع أطول |
| تجميع | — | `torch.compile(module, backend="tpu")` | التجميع المسبق (AOT) باستخدام `TpuBackend` |

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

## الوضع الجشع

لا يتسع نموذج FLUX.1-schnell بالكامل على شريحة v6e واحدة في آنٍ واحد، لذا استخدم [enable_model_cpu_offload()](/docs/diffusers/v0.43.0/ar/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload) لنقل كل نموذج إلى وحدة المعالجة TPU فقط أثناء تشغيله. يقوم هذا الأمر باكتشاف جهاز `"tpu"` تلقائيًا.

```python
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`. يكون الاستدعاء الأول (إحماء) بطيئاً لأنه يقوم بالتجميع؛ أما الاستدعاءات اللاحقة ذات الأشكال نفسها فتعيد استخدام الرسم البياني المُجمَّع.

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

عندما يتسع خط الأنابيب بالكامل على شريحة واحدة، انقله إلى TPU وادمج المحوّل الكامل. يتسع Stable Diffusion 3.5 Medium (حوالي 15GB في bf16) على شريحة v6e واحدة.

```python
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()](/docs/diffusers/v0.43.0/ar/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload)، لا يمكن تتبّع خطافات النقل بواسطة `torch.compile`. قم بتجميع كتل التحوّل المتكررة فقط، باستخدام `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`.

## توازي الموتر

قسّم النماذج التي تتجاوز سعة شريحة واحدة على عدة شرائح. يتجاوز كل من مُشفّر النص (~48GB) والمحوّل (~64GB) في FLUX.2-dev سعة شريحة واحدة، لذا يقوم المثال أدناه بتقسيم كلاهما:

- المحوّل مع [TensorParallelConfig](/docs/diffusers/v0.43.0/ar/api/parallel#diffusers.TensorParallelConfig)، المُمرَّر إلى معامل `parallel_config` في [from_pretrained()](/docs/diffusers/v0.43.0/ar/api/models/overview#diffusers.ModelMixin.from_pretrained). يقرأ كل رتبة شريحتها الخاصة فقط من كل وزن مُجزّأ، لذا لا يتم إنشاء النموذج الكامل أبدًا. للتفاصيل العامة حول TP (`_tp_plan`، colwise/rowwise)، راجع دليل [Tensor parallelism](./tensor_parallelism).
- مشفر النص مع [توازي المصفوفات](https://huggingface.co/docs/transformers/perf_infer_gpu_multi) الخاص بـ Transformers، مع تمرير [DistributedConfig](https://huggingface.co/docs/transformers/v5.22.2/en/expert_parallelism#transformers.DistributedConfig) والشبكة نفسها.

على TPU، قم بتهيئة مجموعة العمليات باستخدام `backend="tpu_dist"` وابنِ الشبكة باستخدام `DeviceMesh("tpu", ...)`.

```python
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 الموجودة على مضيفك.

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

> [!WARNING]
> يستخدم أمر الإقلاع أعلاه وحدة داخلية في TorchTPU (`torch_tpu._internal`)، وليس واجهة برمجة تطبيقات مستقرة، ومن المتوقع أن يتغير في إصدار مستقبلي من TorchTPU. تحقق من [مستودع TorchTPU](https://github.com/google-pytorch/torch_tpu/) لمعرفة الطريقة الحالية لإطلاق المهام متعددة الرقائق.