Google وMeta يقدّمان TorchTPU، واجهة خلفية أصلية لـ PyTorch لوحدات TPU من Google

TorchTPU تعاون مفتوح المصدر بين Google وMeta يقدم واجهة خلفية أصلية لـ PyTorch لوحدات TPU من Google. وهو يتيح تشغيل النماذج وترحيلها إلى TPU دون إعداد إضافي، مع استهداف أداء تنافسي في الاستدلال والتدريب على مستوى القطاع من خلال تحسينات التنفيذ الفوري، وتكامل torch.compile، والتسريع الأصلي للأجهزة.

استُخدمت وحدات TPU من Google لتدريب ونشر نماذج ML كبيرة على نطاق واسع. ومع أن وحدات TPU كانت مرتبطة تاريخيًا بـ TensorFlow وJAX، فقد بدأ هذا الجهد في Google في منتصف 2025 مع فريق PyTorch في Meta. أما تكامل PyTorch/XLA الأقدم، الذي قُدم في 2020 واعتمد على Lazy Tensors، فقد تطلب تعديلات متوغلة في الشفرة المصدرية، وعانى من تأخير حاد عند الترجمة، وافتقر إلى دعم التنفيذ الفوري، وتأخر في العمليات الأولية الموزعة.

مبادئ التصميم وسهولة الاستخدام

بُني TorchTPU كواجهة خلفية TPU قائمة على مبادئ واضحة لـ PyTorch بدلًا من كونه تكييفًا لـ JAX. وتشمل مبادئه الانتماء إلى PyTorch (لا اعتمادات على JAX في وقت التشغيل، ووضع تنفيذ فوري متين، وtorch.compile القائم على Torch Dynamo، ودعم DDP وFSDP وDTensor)، والتكامل السلس مع المنظومة، والمشاركة المجتمعية. التركيز على المستخدمين هو بساطة النقل: تغيير model.to("cuda") إلى model.to("tpu") يجب أن يسمح للنموذج بالعمل دون إعداد إضافي، رغم أن أداء الذروة قد يحتاج إلى ضبط إضافي. كما يسعى المشروع إلى ديناميكية محدودة لتجنب إعادة ترجمة XLA بسبب تغيرات الأشكال، ويستكشف نوى TPU ديناميكية يمكن ترجمتها مسبقًا. ويتطلب استيراد PyTorch على جهاز TPU مع تثبيت حزمة TorchTPU تغييرات طفيفة فقط عبر آلية التحميل التلقائي في PyTorch.

التنفيذ الفوري والترجمة

يوسع TorchTPU موزع PyTorch بواجهة خلفية C++ مسجلة باسم PrivateUse1 ويحدد تخزين موترات TPU عبر آلية at::Tensor القياسية. ويقدم ثلاث وضعيات للتنفيذ الفوري: Strict (افتراضي) يوزع عملية Aten واحدة في كل مرة بشكل غير متزامن؛ Debug (TORCH_TPU_LAUNCH_BLOCKING=1) متزامن مع فحوص إضافية للحدود؛ وDeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) يؤجل ويدمج عمليات متعددة، وتُشير تقارير إلى أنه أسرع بنسبة 50% أو أكثر من Strict eager مع تقليل ذروة ذاكرة HBM. عند نقطة تجسيد مثل استدعاء .to("cpu")، يطبق TorchTPU تحليل الحيوية، ويبني مفتاح ذاكرة مخبئية لبيان العمليات، ويوزع تنفيذًا مخزنًا أو يترجم بشكل غير متزامن. وتكتشف طريقة Repeated Ops Heuristic تسلسلات متكررة، مثل مكدسات كتل transformer، وتؤدي إلى التجسيد المبكر. وتمتد ذاكرة الترجمة المخبئية ثلاثية المستويات عبر ذاكرة محلية للعملية، وذاكرة مشتركة محلية للجهاز، وذاكرة مخبئية دائمة لنظام ملفات بعيد لإعادة الاستخدام عبر المضيفين وإعادة التشغيل.

الأعباء المدعومة والدقة

يدعم TorchTPU حوالي 450 عملية Aten، بما في ذلك المرور الأمامي والخلفي لعمليات Core ATen التي تتطلبها البنى الرئيسية، بالإضافة إلى نوى مخصصة عبر JAX/Pallas أو Helion DSL. وقد تم التحقق منه دون إعداد إضافي على 5,000 من أكثر نماذج Hugging Face تنزيلًا، تغطي LLMs والصوت والكلام والرؤية ومشفرات النصوص. كما تم التحقق من التجزئة القائمة على DTensor، والاستدلال متعدد المضيفين مع vLLM وSGLang، والضبط الدقيق والتدريب مع TorchTitan، والتكامل مع Ray وTorchRec. تستخدم وحدات ضرب المصفوفات في TPU مدخلات bfloat16 مع تراكم FP32، بينما تعمل العمليات التي لا تتضمن انقباضًا في FP32 على الوحدة المتجهة. ويربط TorchTPU واجهات برمجة التطبيقات (APIs) الخاصة بالدقة في PyTorch بهذه الأجهزة: medium، وهو الافتراضي، يستخدم ضرب BF16 مع تراكم FP32؛ وhigh يستخدم خوارزمية من 3 تمريرات بتكلفة حسابية تبلغ نحو 3x؛ وhighest يستخدم 6 تمريرات بعبء إضافي يبلغ نحو 6x. ويؤدي اختيار high أو highest إلى إطلاق تحذير أداء.

المصادر: PyTorch · Google Cloud Blog · Hugging Face Blog