میکسچر آو کیتنز؛ مگاکرنل MoE کرسر
هوش مصنوعیخلاصهٔ کاملتر
کرسر (Cursor) اعلام کرده که مگاکرنل آموزشی MoE خودش به اسم Mixture-of-Kittens یا همون MoK رو اوپنسورس میکنه. به گفتهٔ تیم کرسر، موقع بزرگ کردن آموزش و اجرای Composer، یعنی مدل کدنویسی عاملشون، لایهٔ mixture-of-experts همیشه گلوگاه اصلی بوده و بسته به بار کاری میتونه بیشتر از نصف کل زمان آموزش رو ببلعه.
مشکل اینجا بوده که بهینهسازیهای قبلی فقط بخش محاسباتی رو سریع میکردن و فرض میگرفتن ارتباط بین GPUها جدا هندل میشه؛ ولی توی بار کاری واقعی، همین ارتباط تبدیل به عامل محدودکننده شده بود. رفتن سراغ رکهای GB300 NVL72 هم دو چیز رو عوض کرد: ۷۲ تا GPU توی یک دامنهٔ NVLink قرار میگیرن، و پردازندههای Grace نسبت به GPUها کندن، پس باید کار سمت CPU به حداقل میرسید.
یکی از یافتههای جالب مقاله به انتخاب جهت ارتباط برمیگرده. اغلب پیادهسازیها مثل DeepEP توکنها رو push میکنن، ولی نویسندهها میگن هر کدوم از push و pull جای خودشون رو دارن. با dispatch از نوع pull، جدول زمانبندی خیلی سادهتر میشه و سیگنالدهی بین GPUها کلاً حذف میشه؛ توی بنچمارکهاشون تأخیر سیگنالدهی push حدود ۵.۸ برابر pull بوده.
برای همپوشانی محاسبه و ارتباط، MoK توکنها رو به مینیبچ تقسیم میکنه و اندازهش رو قابل تنظیم گذاشته، چون نه دانهبندی خیلی ریز خوبه نه خیلی درشت. یه بافر حلقوی هم به اسم macrobatch اضافه کردن تا بدون دور ریختن توکن و بدون همگامسازی CPU-GPU با پویایی MoE کنار بیان. کرنل کاملاً قطعیه، یعنی ورودی یکسان همیشه خروجی بیتبهبیت یکسان میده، و از BF16 و MXFP8 پشتیبانی میکنه.
روی بنچمارک تکلایه، MoK تا ۲.۳۷ برابر سریعتر از بهترین بیسلاین توی forward با MXFP8 بوده. مهمتر از اون، توی استک تولید خود کرسر روی ۵۱۲ GPU، تعداد توکن بر ثانیه از حدود ۷۶۱ به ۱۰۷۰ رسیده؛ یعنی ۱.۴۱ برابر. نویسندهها میگن امیدوارن این انتشار، مانع ورود به تحقیقات هوش مصنوعی رو کمتر کنه.
نکات کلیدی:
- MoK کل ارتباط و محاسبات لایهٔ MoE رو توی یک مگاکرنل ادغام میکنه
- ترکیب pull برای dispatch و push برای combine، سیگنالدهی بین GPUها رو حذف میکنه
- بافر حلقوی macrobatch جلوی دور ریختن توکن و همگامسازی CPU-GPU رو میگیره
- خروجی کاملاً قطعیه و هم BF16 و هم MXFP8 پشتیبانی میشن
- سرعت آموزش سرتاسری توی استک تولید کرسر ۱.۴۱ برابر شده
