ThunderKittens من Together AI: لغة جديدة لكتابة GPU kernels بكفاءة

أطلقت Together AI لغة ThunderKittens، وهي لغة برمجة خاصة بمجال محدد (DSL) لكتابة كود محسَّن للعمل على GPU. ويُقدَّم المشروع كمحاولة لتبسيط تطوير الشبكات العصبية، وهو تطوير يتطلب حاليًا معرفة عميقة ببنية العتاد.

لماذا كانت هناك حاجة إلى ذلك

تُعدّ كتابة كود فعّال لـ GPU «سحرًا أسود» بالنسبة لمعظم مهندسي ML. فشرائح مثل NVIDIA H100 تضم tensor cores متخصصة تمنح 94% من إجمالي القدرة الحاسوبية. لكن استخدامها يتطلب الكتابة بـ CUDA، وهي لغة معقدة منخفضة المستوى لا يتقنها سوى قلة في هذه الصناعة.

وفي عام 2023 أظهر الحل الشائع FlashAttention2 أنه يمكن تبسيط العمل مع الانتباه (attention) في المحولات (transformers) بدرجة كبيرة. لكن ذلك مجرد حالة خاصة، إذ يقتصر على خوارزمية لمشغّل واحد بعينه. ويحتاج المطوّرون إلى طريقة عامة لكتابة كود سريع دون جحيم CUDA هذا، طريقة تصلح لأنواع مختلفة من kernels.

كيف تعمل ThunderKittens

تقع ThunderKittens بين نقيضين. فمن جهة هناك CUDA الخالصة: سرعة عالية جدًا لكن تعقيد شديد وحاجز دخول مرتفع. ومن جهة أخرى هناك Triton، الذي يُخفي تفاصيل GPU ويسهّل العمل، لكنه لا يستطيع أحيانًا انتزاع أقصى أداء من العتاد.

وتقترح ThunderKittens طريقًا ثالثًا. فواجهتها البرمجية (API) تشبه PyTorch، وهو ما يألفه كل مطوّر ML اعتاد على العمليات على التنسورات. وفي الوقت نفسه هي شفافة بما يكفي ليفهم المطوّر ما يجري على مستوى العتاد. ويقول المؤلفون: إذا كنت تعرف CUDA، فيمكنك «ترجمة» ThunderKittens في رأسك.

الكائن الأساسي في ThunderKittens هو الـ tile (البلاطة)، وهي مصفوفة بحجم يناسب tensor core. وهذه هي وحدة الحساب التي تتيح تحميل النواة المتخصصة بالكامل واستخراج كل إمكانات العتاد الحديث.

الأداء والنتائج

على الشرائح الحديثة A100 وRTX 4090 تضاهي ThunderKittens سرعة FlashAttention2، مع كود أقصر قليلًا وأوضح. أما على H100 فيتفوق هذا الحل: فهو أسرع من FlashAttention2 في المرور الأمامي وفي الانتشار الخلفي (backpropagation) على حد سواء. أي أنه لا توجد مقايضة بين نظافة الكود والسرعة.

وقد كتب المؤلفون بالفعل بـ ThunderKittens عدة kernels لخوارزميات أخرى:

• Based — نسخة محسَّنة من الانتباه الخطي
• Hedgehog وkernels متخصصة أخرى للمحولات
• عدة حلول تتفوق في السرعة بشكل لافت على نسخ Triton

ومن الملاحظات المثيرة للاهتمام أن مهندسين أنهوا دورة CUDA مدتها ساعتان فقط تمكنوا من كتابة كودهم الخاص بـ ThunderKittens. ويدل ذلك على أن اللغة تبسّط التطوير فعلًا، وإن كانت لا تُخفي تفاصيل العتاد.

مشروع مفتوح وتعليم

يقول مؤلفو Together AI بصراحة: هذا art project (مشروع فني). فلا تتوقعوا تحديثات منتظمة ولا معالجة كل شكاواكم في متتبّع الأخطاء. وقد أُتيح المشروع للجميع لأن المطوّرين يرون أن مشاركة الأفكار والأدوات مع المجتمع أمر ممتع.

وإلى جانب ThunderKittens أُطلق NanoGPT-TK، وهو نسخة من مشروع NanoGPT الشهير لأندريه كارباثي، أُعيدت فيها كتابة نوى الحساب الرئيسية بـ ThunderKittens. وقد جرى ذلك خصيصًا لأغراض التعليم وإظهار الإمكانات. ويُعدّ NanoGPT منذ زمن من أفضل المشاريع في مجال AI لفهم كيفية عمل تدريب المحولات من الصفر.

ماذا يعني ذلك

تُظهر ThunderKittens أن في مجال AI فجوة حقيقية بين سهولة التجريدات (PyTorch وTriton) والتحكم في العتاد الفعلي (CUDA). ويبدو أن المطوّرين مستعدون لكتابة كود أكثر تعقيدًا بقليل إذا منحهم ذلك التحكم وسرعة حقيقية في الممارسة.

وبالنسبة لمهندسي ML، قد يعني ذلك أن نقل نموذج مدرَّب من شريحة إلى أخرى سيصبح أسهل مستقبلًا، إذ يكفي إعادة كتابة بضع GPU kernels بدلًا من إعادة بناء نصف البنية التحتية. أما بالنسبة للباحثين فهي أداة للتجريب السريع مع خوارزميات متخصصة، ولا تتطلب دورة CUDA تستغرق شهرًا.