هذا ما كنا نستخدمه خلال الأشهر التسعة الماضية: اجعل تدريب MoEs أسرع ~ ضعف وذاكرة أقل ~2x! أبرز النقاط التاريخية: - عادة ما تستغرق MoE أكبر قدر من الوقت والذاكرة في النماذج الحديثة. اتضح أنه يمكن إعادة كتابة مرور MoE للخلف رياضيا لتقليل ذاكرة التفعيل التي تحتاج لتخزينها في المحرك الأمامي بمقدار ~2x، مما يؤدي إلى نفس التدرجات دون الحاجة لإعادة حساب matmul إضافية. أعجبني هذا النتيجة كثيرا، لأنها تجمع بين الرؤى الخوارزمية والنظامية. - تحليل عنق الزجاجة في طبقة MoE يؤدي إلى استراتيجية تحسين طبيعية: تقليل قراءات/كتابة الذكريات قدر الإمكان! جمع المدخلات للدفع الأمامي ودرجة الإخراج لاستخدام التردد قد يستغرق أحيانا نفس الوقت الذي تستغرقه وحدات GEMM المجمعة. نقوم بدمج الجمع مع GEMM + التداخل المدمج للوصول إلى الذاكرة وحساب لجعل الطبقة بأكملها أسرع ~2x. - حساب top-k للتوجيه الخبير قد يستغرق وقتا طويلا بشكل مفاجئ، ~15-20٪ من طبقة MoE بأكملها! يستخدم ال top k القياسي الجزئية العليا (Radix top-k algo)، وهو ممتاز لل k الكبير لكنه غير مثالي لل k الصغير. أعدنا كتابة top-k باستخدام خوارزمية top-k الثنائية (bitonic top-k)، وأحيانا تكون أسرع ب 20-30 مرة من top-k في بايتورش! جميع النوى الرئيسية مكتوبة بلغة Cute-DSL، لذا يجب أن يكون من السهل توسيعها (وتركيبها :D). حبوب الهوبر خرجت، وحبوب بلاكويل جاهزة تقريبا. كانت نماذج MoE أقل كفاءة في الأجهزة بمرتين للتدريب، ونأمل أن يغير Sonic-MOE ذلك.