‏Olmo-core 3؛ زیرساخت باز Ai2 برای آموزش مدل‌های MoE غول‌پیکر

Ai2 تو این پست از Olmo-core 3 رونمایی کرده؛ نسخه‌ی جدید فریم‌ورکی که مدل‌های Olmo باهاش آموزش می‌بینن. محور اصلیش یه سیستم بازطراحی‌شده برای آموزش مدل‌های MoE هست (یعنی Mixture of Experts، مدلی که از چندین زیرشبکه‌ی تخصصی یا expert ساخته شده و برای هر توکن فقط چندتاشون کار می‌کنن). نسل بعدی Olmo هم قراره روی همین معماری ساخته بشه.

مشکل اینجاست که حتی اگه هر توکن فقط بخشی از مدل رو استفاده کنه، کل مدل باید تو حافظه‌ی GPUها باشه. فرستادن توکن‌ها به expert درست روی یه کلاستر هم هزینه‌ی ارتباطی داره. به گفته‌ی Ai2، با بزرگ شدن مدل این هزینه‌ها می‌تونه بیشتر مزیت MoE رو بخوره. تو یه آزمایش، تعداد expertها رو از ۸ به ۱۲۸ رسوندن و پارامترهای فعال هر توکن رو حدود 3.2B نگه داشتن. کل پارامترها از 4.6B به 47B رسید و سرعت آموزش کمتر از ۵٪ افت کرد.

تغییر اصلی اینه که از FSDP (روشی که وزن‌ها رو برای هر batch بین GPUها جمع و دوباره پخش می‌کنه) به DDP رفتن. حالا expertها ثابت روی GPUها می‌مونن و داده به سمتشون فرستاده می‌شه. روی هشت GPU مدل B300، یه MoE با ۴۷ میلیارد پارامتر به ۵۲ هزار توکن در ثانیه برای هر GPU رسید، در برابر ۱۹٬۴۰۰ تا با پیاده‌سازی قبلی؛ یعنی حدود ۲.۷ برابر.

برای پخش مدل روی سخت‌افزار سه تکنیک با هم کار می‌کنن: expert parallelism که expertها رو بین GPUها تقسیم می‌کنه، pipeline parallelism که لایه‌ها رو بین گروه‌های GPU پخش می‌کنه، و یه optimizer توزیع‌شده که داده‌ی آپدیت وزن‌ها رو روی همه‌ی GPUها کپی نمی‌کنه. کنارش چند بهینه‌سازی هم هست: routing که متادیتاش روی GPU می‌مونه تا CPU معطل نشه، و grouped GEMM که کلی محاسبه‌ی کوچیک expertها رو یکجا اجرا می‌کنه.

از فرمت کم‌دقت MXFP8 هم پشتیبانی می‌کنه. روی چهار GPU مدل B300، سرعت آموزش حدود ۲۱٪ بیشتر از BF16 شد و اوج حافظه از 103GiB به 95GiB رسید. بزرگ‌ترین تستشون یه مدل ۱.۲ تریلیون پارامتری روی ۵۱۲ GPU بوده که به 858 TFLOP/s برای هر GPU رسیده. البته با routing تصادفی سنجیده شده، نه کیفیت یه مدل واقعاً آموزش‌دیده. با DeepEP v2 هم یه تست کوتاه تا ۲.۳۸ تریلیون پارامتر رفتن.

گزارش فنی چند تا یافته‌ی جالب هم داره. مثلاً امتیازی که قرار بود توزیع متعادل بین expertها رو تشویق کنه، گاهی بهتر می‌شد در حالی که بار واقعی نامتعادل‌تر شده بود. اسمش رو گذاشتن token gerrymandering. یافته‌ی دیگه اینه که هم‌پوشانی ارتباط و محاسبه همیشه سرعت رو بیشتر نمی‌کنه و گاهی کندترش هم کرده.

نکات کلیدی:

• رفتن از FSDP به DDP؛ حدود ۲.۷ برابر throughput روی یه MoE ۴۷ میلیاردی
• افزایش expertها از ۸ به ۱۲۸ با کمتر از ۵٪ افت سرعت آموزش
• MXFP8 حدود ۲۱٪ سریع‌تر از BF16 و ۸GiB حافظه‌ی کمتر
• تست ۱.۲ تریلیون پارامتری روی ۵۱۲ GPU با 858 TFLOP/s برای هر GPU
• کاملاً متن‌باز و پایه‌ی نسل بعدی Olmo