# Google TPU

يعمل [Google Cloud TPUs](https://cloud.google.com/tpu) مع PyTorch بشكل أصلي من خلال حزمة [torch_tpu](https://github.com/google-pytorch/torch_tpu)، التي تسجّل جهاز `tpu`. تُجمّع `torch_tpu` العمليات التي تُنفّذها لكل شكل إدخال تراه، لذا فإن إبقاء أشكال الإدخال ثابتة يتجنّب إعادة التجميع.

## التثبيت

قم بتثبيت `torch_tpu` باتباع التعليمات الموجودة في [مستودع torch_tpu](https://github.com/google-pytorch/torch_tpu). يتم تحميل `torch_tpu` تلقائياً عند استيراد `torch`. تأكد من أن PyTorch يرى شرائح TPU.

```py
import torch

print(torch.accelerator.is_available(), torch.accelerator.device_count())
```

يطبع `True` وعدد شرائح TPU، على سبيل المثال `True 4`.

## تشغيل الاستدلال

`device_map="auto"` لا يكتشف وحدات TPU. يعتمد التوزيع التلقائي على [Accelerate](https://huggingface.co/docs/accelerate/index) لذكر الأجهزة المتاحة، لكن Accelerate لا يذكر أجهزة TPU، لذا يتم تحميل النموذج على وحدة المعالجة المركزية (CPU) دون أي تنبيه. مرّر `device_map="tpu"` إلى [from_pretrained()](/docs/transformers/main/ar/main_classes/model#transformers.PreTrainedModel.from_pretrained) بشكل صريح، أو مرّر `device="tpu"` إلى [Pipeline](/docs/transformers/main/ar/main_classes/pipelines#transformers.Pipeline).

مرّر مدخلات التوليد على وحدة المعالجة المركزية (CPU) بدلاً من استدعاء `.to(model.device)`. تنقل [generate()](/docs/transformers/main/ar/main_classes/text_generation#transformers.GenerationMixin.generate) فقط المصفوفات التي تستهلكها `forward` إلى وحدة المعالجة المخصصة (TPU). تبقى المصفوفات التي تنمو في كل خطوة (تسلسل الإخراج، `attention_mask`) على وحدة المعالجة المركزية، مما يحافظ على استقرار الرسم المُجمّع. تُرجع `generate` الإخراج على نفس الجهاز الذي توجد عليه المدخلات. راجع [إبقاء مصفوفات التوليد على وحدة المعالجة المركزية](./kv_cache#keep-generation-tensors-on-cpu) لمزيد من التفاصيل. تستخدم النماذج SDPA افتراضياً على وحدة المعالجة المخصصة.

```py
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="tpu"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=True, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))
```

```txt
The French Bread Law states that a farmer can sell their harvest of bread at a certain price. The price is determined by the market price and the price of the bread, and the price is set by the farmers' decision-making committee. What are the implications of this law on
```

مع ذاكرة مؤقتة ثابتة، يجمع [generate()](/docs/transformers/main/ar/main_classes/text_generation#transformers.GenerationMixin.generate) مسار التقدّم باستخدام `torch.compile` وخلفية `tpu` تلقائيًا. مرّر `disable_compile=True` إلى `generate` للاستثناء من ذلك.

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

`torch_tpu` يشغّل عملية واحدة لكل شريحة TPU، وترى كل عملية شريحتها على أنها `tpu:0`. أنشئ [DistributedConfig](/docs/transformers/main/ar/expert_parallelism#transformers.DistributedConfig) بحجم التوازي على مستوى الموتر، وقم بتمريره إلى [from_pretrained()](/docs/transformers/main/ar/main_classes/model#transformers.PreTrainedModel.from_pretrained). يقوم Transformers بتجزئة النموذج عند وقت التحميل.

```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=4),
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
if torch.distributed.get_rank() == 0:
    print(tokenizer.decode(output[0], skip_special_tokens=True))
```

يتطلب `torch_tpu` ضبط بعض متغيرات البيئة قبل بدء العمليات، وهو ما لا يقوم به `torchrun` تلقائياً. استدعِ `torch.tpu.distributed.set_tpu_launch_env()` ثم ابدأ `torchrun` من Python.

```py
# launch.py
import sys

import torch
from torch.distributed.run import main

torch.tpu.distributed.set_tpu_launch_env()
main(sys.argv[1:])
```

مرّر حُجج `torchrun` المعتادة إلى مُشغِّل التشغيل.

```bash
python launch.py --nproc_per_node=4 generate.py
```

تتطلب بعض النماذج تكييفًا قبل أن يمكن تقسيمها على شرائح TPU.

لا يمكن تقسيم طبقات الانتباه في النماذج التي لا يكون عدد رؤوس المفاتيح والقيم فيها قابلاً للقسمة على `tp_size`، مما يؤدي إلى فشل في مسار الأمامي (forward pass) بسبب خطأ في الأبعاد. بدلاً من ذلك، قم بتكرار إسقاطات الانتباه على كل شريحة عن طريق تجاوز قواعدها في `tp_plan`. القواعد التي لا تقوم بتجاوزها ستحتفظ بالخطة الافتراضية للنموذج. على سبيل المثال، يحتوي [GLM-4-9B-0414](https://huggingface.co/zai-org/GLM-4-9B-0414) على رؤوس مفاتيح وقيم عددها 2، ولا يمكن تقسيمها عبر 4 شرائح.

```py
model = AutoModelForCausalLM.from_pretrained(
    "zai-org/GLM-4-9B-0414",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(
        tp_size=4,
        tp_plan={
            "layers.*.self_attn.q_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.k_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.v_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.o_proj": "replicated_with_grad_allreduce",
        },
    ),
)
```

تفشل النماذج ذات الأوزان المرتبطة، مثل Qwen3-0.6B، في التحميل عندما يكون `tp_size` أكبر من 1.

## استخدام نواة Hub

يدعم مكتبة [kernels](https://github.com/huggingface/kernels) معالجات TPU منذ الإصدار v0.17. تُغلّف نوى TPU كود JAX و[Pallas](https://docs.jax.dev/en/latest/pallas/index.html)، لذا ثبّت `jax` إلى جانب `kernels`.

```bash
pip install "transformers[kernels]" jax
```

راجع [تحميل النوى](./kernel_doc/loading_kernels) لمعرفة كيفية تحميل نواة من Hub.

## القيود

- `device_map="auto"` يحمّل النموذج على وحدة المعالجة المركزية (CPU) بدلاً من وحدة المعالجة المخصصة (TPU). مرّر `device_map="tpu"` بشكل صريح.
- تُحفّز الأشكال الديناميكية إعادة التجميع. يتباطأ توليد النص عندما تتغير أشكال المدخلات من خطوة إلى أخرى. استخدم ذاكرة مؤقتة ثابتة لتسريع عملية التوليد.
- يمكن لكل عملية الوصول إلى شريحة TPU واحدة فقط. استخدم عملية واحدة لكل شريحة للتشغيل على عدة شرائح.
- لا تزال تكاملات Transformers للتدريب والتكميم على TPU قيد التطوير، لذا قد لا تعمل بعض الميزات بعد. يرجى [فتح مشكلة](https://github.com/huggingface/transformers/issues) إذا واجهت مشكلة.

