ThunderKittens من Together AI: لغة جديدة لكتابة نوى GPU فعالة

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

لماذا كان هذا ضروريًا

كتابة أكواد فعالة لـ GPU تبدو وكأنها سحر أسود لمعظم مهندسي التعلم الآلي. على المعالجات مثل NVIDIA H100، توجد نوى موتّر متخصصة توفر 94% من إجمالي قوة الحوسبة. لكن لاستخدام هذه النوى، يجب الكتابة بلغة CUDA - لغة معقدة جداً في المستوى المنخفض، لا يفهمها سوى عدد قليل من الخبراء في الصناعة.

في عام 2023، أظهر الحل الشهير FlashAttention2 أنه من الممكن تبسيط العمل مع آلية الانتباه في محولات Transformer بشكل كبير. لكن هذا كان حالة خاصة فقط - خوارزمية لمشغل واحد محدد. يحتاج المطورون إلى طريقة عامة لكتابة أكواد سريعة دون الحاجة إلى CUDA المعقد هذا، وتعمل على نوى مختلفة.

كيف يعمل ThunderKittens

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

يقدم ThunderKittens المسار الثالث. واجهة البرمجة الخاصة به تشبه PyTorch، وهي مألوفة لجميع مطوري التعلم الآلي الذين اعتادوا على العمليات الموتّرية. وفي الوقت نفسه، فهو شفاف بما يكفي بحيث يفهم المطور ما يحدث على مستوى الأجهزة. يقول المطورون: إذا كنت تعرف CUDA، يمكنك "ترجمة" ThunderKittens في رأسك.

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

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

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

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

• Based - نسخة محسّنة من الانتباه الخطي
• Hedgehog وغيرها من النوى المتخصصة لـ Transformers
• عدة حلول تتميز بسرعة أفضل عن نسخ Triton

نقطة مثيرة للاهتمام: المهندسون الذين أكملوا فقط دورة تدريبية قصيرة مدتها ساعتان عن CUDA تمكنوا من كتابة أكوادهم الخاصة على ThunderKittens. هذا يشير إلى أن اللغة بالفعل تبسط التطوير، رغم عدم إخفائها لتفاصيل الأجهزة.

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

يصرح المطورون من Together AI بصراحة: هذا مشروع فني. لا تتوقع تحديثات منتظمة ودعماً لكل شكاويك في متتبع الأخطاء. تم إطلاق المشروع بشكل علني لأن المطورين يعتقدون أنه من المثير للاهتمام مشاركة الأفكار والأدوات مع المجتمع.

جنباً إلى جنب مع ThunderKittens، أطلقوا NanoGPT-TK - نسخة من مشروع NanoGPT الشهير من Andrej Karpathy، حيث تمت إعادة كتابة أساسيات حساب النوى على ThunderKittens. تم القيام بهذا بشكل خاص للتعليم والتوضيح.

اعترف NanoGPT منذ فترة طويلة بأنه أحد أفضل المشاريع في الذكاء الاصطناعي لفهم كيفية عمل تدريب Transformers من الصفر.

ما معنى هذا

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

بالنسبة لمهندسي التعلم الآلي، قد يعني هذا أنه في المستقبل، نقل النموذج المدرب من معالج واحد إلى آخر سيكون أسهل - يكفي إعادة كتابة عدد قليل من نوى GPU بدلاً من إعادة بناء أنصاف البنية الأساسية.

بالنسبة للباحثين، هذه أداة للتجريب السريع مع الخوارزميات المتخصصة التي لا تتطلب دورة تدريبية مدتها شهر واحد في CUDA.