إذا حاولت يوماً تشغيل PyTorch على وحدة TPU، فمن المرجح أنك نقلت شيفرتك إلى torch_xla، أو انتهى بك الأمر إلى تعلم JAX. مع TorchTPU، وهو خلفية أصلية لـ PyTorch، أصبحت TPU الآن مجرد جهاز آخر. هذا يفتح الدعم الأصلي لوحدات TPU في منظومة Hugging Face. تعمل Transformers وDiffusers وKernels بشكل أصلي على TPU. اختر نموذجاً من الـ Hub، وأرسله إلى "tpu"، وينتهي الأمر.
مهلاً، ما هي TPU؟ 🤔
TPU هي اختصار لـ Tensor Processing Unit(وحدة المعالجة الموترية). وحدات TPU هي مسرّعات الذكاء الاصطناعي الخاصة بشركة Google. إنها رقائق بُنيت لوظيفة واحدة: العمليات الحسابية المصفوفية التي تقضي الشبكات العصبية معظم وقتها فيها. صمّمت Google هذه الوحدات لمدة نحو عشر سنوات، وتستخدمها لتدريب نماذجها الخاصة وتقديمها، بما في ذلك Gemini.
تتوفر وحدات TPU عبر Google Cloud، من رقاقة واحدة وحتى عناقيد (pods) تضم آلاف الرقائق. تتألق عندما تحتاج إلى قدر كبير من الحوسبة، سواء للتدريب أو للتقديم على نطاق واسع.
لكن العقبة كانت دائماً في البرمجيات. نشأت منظومة TPU حول TensorFlow وJAX، واضطر مستخدمو PyTorch إلى بذل جهد أكبر للوصول إليها.
الطريقة القديمة
حتى الآن، كان تشغيل PyTorch على وحدات TPU يتطلب المرور عبر PyTorch/XLA (torch_xla). كان يعمل، لكن كان له طابعه الغريبة. كانت الموترات (tensors) كسولة (الموترات، وليست أنت)، فلم تكن شيفرتك تنفَّذ دائماً عندما تظن أنها تُنفَّذ. وكانت التغييرات الصغيرة قد تؤدي إلى إعادة تجميع الرسم البياني، وكان تصحيح الأخطاء صعب الفهم.
لتسهيل الحياة، بنينا Optimum TPU. وفّر هذا المشروع دعم الاستدلال والتدريب لنماذج Hugging Face على وحدات TPU. والآن بعد أن يوفر TorchTPU دعماً أصلياً، لم يعد Optimum TPU مطلوباً لتشغيل نماذج Hugging Face مع Transformers وDiffusers.
أهلاً بك في TorchTPU 🚀
TorchTPU يضيف TPU كخلفية جديدة إلى PyTorch. تكتب شيفرات PyTorch الاعتيادية. يأتي مع نوى Pallas مضبوطة لوحدات TPU، ودعماً للتدريب الموزع، وتكاملاً مع torch.compile للأوقات التي تريد فيها سرعة أكبر. تثبيته على جهاز TPU VM يتطلب أمراً واحداً:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
يمنحك الفهرس الإضافي نسخة CPU من torch. بدونه، يسحب pip نسخة CUDA الأكبر بكثير، والتي لا تحتاجها على TPU.
مكتباتك المفضلة، الآن على TPU
Transformers
بُنيت Transformers الإصدار v5 حول فكرة واحدة: أن تكون المرجع الأساسي لتعريفات النماذج، حتى تتمكن من العمل في أي مكان. و"أي مكان" تشمل الآن وحدات TPU. لتحميل نموذج على TPU، مرّر device_map="tpu". هذا كل ما في الأمر.
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="tpu", # will load model directly on the TPU
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu")
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))
يعمل النموذج في الوضع السريع (eager mode) افتراضياً. يمكنك أيضاً استخدام torch.compile للحصول على أداء أفضل، أو تشغيله في بيئة موزعة.
صور؟ فيديوهات؟ Diffusers يرسمها على وحدات TPU 🎨
أصبحت TPU الآن جهازاً من الدرجة الأولى في Diffusers. للنماذج الأولية أو تصحيح الأخطاء، لا يحتاج وضع eager سوى pipe.to("tpu"). وللنشر، يضيف الوضع المُجمَّع استدعاء compile واحداً. تبقى خطوات الأنابيب والشيفرة لديك كما هي، ولا حاجة للحفاظ على نسخة TPU منفصلة.
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu") # eager mode: that's the whole change
# optional: compiled mode
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0]
الاستدعاء الأول يقوم بالتجميع، وكل استدعاء لاحق بنفس حجم الصورة يعيد استخدام الرسم البياني.
هل النماذج كبيرة جداً بحيث لا تتسع لشريحة واحدة؟ مرِّر TensorParallelConfig إلى from_pretrained وسيتم تقسيم المحوّل (transformer) عبر الشرائح أثناء تحميله، فلا يلزم أبداً أن يتسع النموذج الكامل في مكان واحد. تدعمه خطوات أنابيب شهيرة كثيرة مباشرةً دون إعداد، والتي لا تدعمه لا تحتاج سوى tp_plan . الدليل الكامل موجود في دليل Diffusers لـ TPU .
إليك بعض الأمثلة المُشغَّلة على Cloud TPU v6e-4:
TP=4 (كل نموذج موزَّع عبر أربع شرائح) · bf16 · حجم الدفعة 1
| MiniMax-H3 (فيديو + صوت) | Qwen-Image-2.1 | |
|---|---|---|
| "دمية هاغين فيس المحشوّة الناعمة من الصورة، وجه أصفر دائري مبتسم مع وجنتين وردية ويدّين صغيرتين مفتوحتين، تقف في مقدمة سفينة ترييمة يونانية قديمة [...] صوت أمواج متلاطمة، وخشب يتقرّع، وريح في الحبال، ونورس بعيد، وتصاعد موسيقي أوركسترالي ملحمي هادئ." · 1344×768، 5.2 ثانية، 49 خطوة · 188 ثانية مُجمَّع، 262 ثانية eager | "كابيبارا يرتدي قبعة ساحر، يقرأ كتاباً على ضوء الشمعة، لوحة زيتية" · 1024×1024، 40 خطوة · 3.5 ثانية مُجمَّع، 8.0 ثانية eager | "لافتة متجر نيون مكتوب عليها 'QWEN IMAGE 2.1'، ليلة ممطرة، انعكاسات على رصيف مبلّل" · 2048×2048، 40 خطوة |
النوى (Kernels)
يحتوي مركز Kernels على نوى TPU جاهزة للاستخدام. حمِّل إحداها عبر مكتبة Kernels وامنح نموذجك دفعة سرعة.
تُكتب نوى TPU بلغة Pallas، مما يمنحك تحكمًا منخفض المستوى في العتاد. ويمكنك أيضًا استيراد نوى من مكتبة Tokamax . إليك نواة ragged-dot من Tokamax مدمجة في نموذج Mixture-of-Experts:
from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, version=1, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-30B-A3B-Instruct-2507",
experts_implementation="tokamax_ragged_dot",
)
كلما كبر حجم الدفعة، زادت المكسب. عند حجم دفعة 64، يكون فك الترميز أسرع بمقدار 21.5× على Qwen3-30B-A3B وأسرع بمقدار 12.8× على Gemma 4 26B-A4B.
إنتاجية فك الترميز على TPU v6e-8، نواة Hub مقابل التنفيذ الافتراضي.
نحن نجّربه حتى لا تضطر أنت لذلك 🧪
نعمل باستمرار على تحسين دعم TPU في transformers من خلال تشغيل حزمة اختبارات المكتبة. ونعمل على إضافة اختبار CI ليلي على TPU وجعل النتائج متاحة في (transformers dashboard)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard].
لوحة CI الخاصة بـ Transformers، مع نتائج TPU بجوار مسرّعات أخرى. عمل قيد التنفيذ!
الأمر ليس أخضر بالكامل بعد، ونحن مرتاحون لذلك. من الاختبارات التي تعمل على TPU، ينجح 96.5%. وكل فشل نكتشفه يتحول إلى سكربت إعادة إنتاج صغير ومشكلة (issue) في مستودع TorchTPU. حتى الآن أبلغنا عن 34 مشكلة، وتم إصلاح 19 منها بالفعل. كما كشف التشغيل على TPU بعض الأخطاء في Transformers نفسه، مثل أخطاء في StaticCache وخطط Qwen3 للموازاة الشريطية للموترات (tensor-parallel). أصلحناها في المصدر الأصلي.
إذا تعطل شيء ما على TPU، فإننا نريد أن نكتشف ذلك قبلك.
ابدأ الآن
لم تعد TPUs مهمة جانبية لمستخدمي PyTorch. ثبّت TorchTPU، التقط نموذجًا من الـ Hub، وأرسله إلى "tpu".
لا نطيق صبرنا لرؤية ما ستبنيه 🤗. إذا واجهت مشكلة، افتح issue وأخبرنا بها.