Google وMeta تفتحان المصدر لـ TorchTPU لدعم PyTorch على وحدات TPU

TorchTPU هو مشروع مفتوح المصدر طورته Google بالتعاون مع Meta، صُمم لتوفير خلفية PyTorch أصلية لوحدات Google TPU. تتيح هذه المبادرة تشغيل أحمال عمل PyTorch على عتاد TPU بأقل تعديلات في الشيفرة، وتدعم التنفيذ الفوري والمُجمّع، والحوسبة الموزعة، والتحسينات الخاصة بالعتاد.

بدأت Google في تسهيل تشغيل أحمال عمل PyTorch على وحدات TPU في 2020 عبر PyTorch/XLA، الذي استخدم JAX وXLA وlazy tensors. لكن Google رصدت قيوداً كبيرة في هذا النهج، منها الحاجة إلى تعديلات جوهرية في الشيفرة المصدرية، وزمن تأخير حاد في التجميع، ودعم غير كافٍ لنمط التنفيذ الفوري الأصلي في PyTorch.

يعتمد TorchTPU استراتيجية مختلفة جذرياً، إذ يعطي الأولوية للتنفيذ الفوري ويسمح لأحمال عمل PyTorch الحالية بالعمل بتعديلات طفيفة. ويتطلب الانتقال من وحدات GPU إلى TPU ببساطة استبدال جهاز `cuda` بـ `tpu` وخلفية التوزيع `nccl` بـ `tpu_dist`. ويمكن للمطورين أيضاً استخدام واجهة `torch.accelerator.current_accelerator` الجديدة لتحديد المسرّع المتاح.

يدعم TorchTPU ثلاثة أنماط للتنفيذ الفوري: strict وdebug وDeferAndFuse. ونمط DeferAndFuse هو نهج تنفيذ هجين بين الفوري والمُجمّع، وتقول Google إنه قد يكون أسرع بنسبة 50% أو أكثر من Strict Eager مع تقليل ذروة استخدام ذاكرة HBM.

لتقييم قدرة TorchTPU على تشغيل نماذج PyTorch الحالية دون تعديلات، أنشأت Google خط تكامل قيّم نحو 5000 نموذج من Hugging Face عبر أكثر من 500 بنية في النمطين الفوري والمُجمّع على أجيال مختلفة من عتاد TPU.

للتحقق من الأداء، اختبرت Google نموذج Muse Glimmer ذا 30 مليار معامل والمُنفّذ بـ PyTorch القياسي. على شريحة TPUv7 واحدة، بلغت سرعة الاستدلال 25.7k token/s عند طول تسلسل 5120 و22.4k token/s عند 10,240. أما للتدريب، فباستخدام أربع شرائح TPUv7، بلغت السرعة 5.6k token/s و4.6k token/s عند طولي التسلسل نفسيهما.

أجرت Google أيضاً تجارب على نموذج Qwen3 MoE 235B لأحمال عمل التدريب، مع التوسع من 1024 إلى 8192 شريحة TPUv7، وأبلغت عن مقاييس أداء تنافسية.

المصدر: InfoQ