Google TPU

يعمل Google Cloud TPUs مع PyTorch بشكل أصلي من خلال حزمة torch_tpu، التي تسجّل جهاز tpu. تُجمّع torch_tpu العمليات التي تُنفّذها لكل شكل إدخال تراه، لذا فإن إبقاء أشكال الإدخال ثابتة يتجنّب إعادة التجميع.

التثبيت

قم بتثبيت torch_tpu باتباع التعليمات الموجودة في مستودع torch_tpu. يتم تحميل torch_tpu تلقائياً عند استيراد torch. تأكد من أن PyTorch يرى شرائح TPU.

import torch

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

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

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

device_map="auto" لا يكشف عن وحدات TPU. يعتمد التوزيع التلقائي على Accelerate لذكر الأجهزة المتاحة، ولا يذكر Accelerate أجهزة TPU، لذا يتم تحميل النموذج على وحدة المعالجة المركزية (CPU) بصمت. مرّر device_map="tpu" إلى from_pretrained() بشكل صريح، أو مرّر device="tpu" إلى Pipeline.

مرّر مدخلات التوليد على وحدة المعالجة المركزية (CPU) بدلاً من استدعاء .to(model.device). تنقل generate() فقط المصفوفات التي تستهلكها forward إلى وحدة المعالجة المخصصة (TPU). تبقى المصفوفات التي تنمو في كل خطوة (تسلسل الإخراج، attention_mask) على وحدة المعالجة المركزية، مما يحافظ على استقرار الرسم المُجمّع. تُرجع generate الإخراج على نفس الجهاز الذي توجد عليه المدخلات. راجع إبقاء مصفوفات التوليد على وحدة المعالجة المركزية لمزيد من التفاصيل. تستخدم النماذج SDPA افتراضياً على وحدة المعالجة المخصصة.

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

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() مسار التقدّم باستخدام torch.compile وخلفية tpu تلقائيًا. مرّر disable_compile=True إلى generate للاستثناء من ذلك.

توازي الموتر

torch_tpu يشغّل عملية واحدة لكل شريحة TPU، وترى كل عملية شريحتها على أنها tpu:0. أنشئ DistributedConfig بحجم التوازي على مستوى الموتر، وقم بتمريره إلى from_pretrained(). يقوم Transformers بتجزئة النموذج عند وقت التحميل.

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.

# launch.py
import sys

import torch
from torch.distributed.run import main

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

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

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

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

لا يمكن تقسيم طبقات الانتباه في النماذج التي لا يكون عدد رؤوس المفاتيح والقيم فيها قابلاً للقسمة على tp_size، مما يؤدي إلى فشل في مسار الأمامي مع خطأ في الشكل. بدلاً من ذلك، انسخ إسقاطات الانتباه على كل شريحة عن طريق تجاوز قواعدها في tp_plan. القواعد التي لا تتجاوزها تحتفظ بالخطة الافتراضية للنموذج. على سبيل المثال، يحتوي GLM-4-9B-0414 على رؤوس مفاتيح وقيم عددها 2، ولا يمكن تقسيمها عبر 4 شرائح.

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 معالجات TPU منذ الإصدار v0.17. تُغلّف نوى TPU كود JAX وPallas، لذا ثبّت jax إلى جانب kernels.

pip install "transformers[kernels]" jax

راجع تحميل النوى لمعرفة كيفية تحميل نواة من Hub.

القيود

• device_map="auto" يحمّل النموذج على وحدة المعالجة المركزية (CPU) بدلاً من وحدة المعالجة المخصصة (TPU). مرّر device_map="tpu" بشكل صريح.
• تُحفّز الأشكال الديناميكية إعادة التجميع. يتباطأ توليد النص عندما تتغير أشكال المدخلات من خطوة إلى أخرى. استخدم ذاكرة مؤقتة ثابتة لتسريع عملية التوليد.
• يمكن لكل عملية الوصول إلى شريحة TPU واحدة فقط. استخدم عملية واحدة لكل شريحة للتشغيل على عدة شرائح.
• لا تزال تكاملات Transformers للتدريب والتكميم على TPU قيد التطوير، لذا قد لا تعمل بعض الميزات بعد. يرجى فتح مشكلة إذا واجهت مشكلة.