جهش ۲.۷ برابری در آموزش غول‌های هوش مصنوعی؛ مؤسسه Ai2 از OLMo-Core 3 رونمایی کرد

بازطراحی پشته آموزشی اولمو توسط Ai2 برای مدل‌های MoE تریلیون پارامتری

مؤسسه هوش مصنوعی آلن (Ai2) به‌تازگی از OLMo-core 3 پرده‌برداری کرده است؛ یک پشته و زیرساخت آموزشی کاملاً متن‌باز که برای مدل‌های غول‌پیکر با معماری ترکیبی از متخصصان (MoE) توسعه یافته است. این مؤسسه اعلام کرده که قصد دارد از همین زیرساخت برای آموزش نسل بعدی مدل‌های پرچمدار اولمو (Olmo) استفاده کند.

مدل‌های ترکیبی از متخصصان (MoE) شبکه‌های عصبی فیدفوروارد متعددی در دل خود دارند و یک روتر هوشمند، هر توکن را فقط به بخش کوچکی از این متخصص‌ها هدایت می‌کند. این فعال‌سازی پراکنده باعث می‌شود ظرفیت پارامترهای مدل سر به فلک بکشد بدون اینکه بار محاسباتی به همان اندازه سنگین شود. با این حال در مقیاس کلاسترهای محاسباتی بزرگ، چالش‌هایی مثل تبادل توکن‌ها، جابجایی وزن‌ها، مدیریت وضعیت بهینه‌ساز (optimizer state) و هماهنگی مسیریابی می‌توانند تمام این صرفه‌جویی را بر باد دهند. حالا پشته OLMo-core 3 دقیقاً طوری طراحی شده که حتی با بزرگ‌تر شدن استخر متخصصان، توان پردازشی و سرعت تولید توکن تقریباً ثابت و پایدار بماند.

جهش ۲.۷ برابری سرعت روی هشت پردازنده گرافیکی B300 انویدیا

مؤسسه Ai2 نتایج دو بنچمارک اولیه را منتشر کرده که به‌خوبی نشان می‌دهند این پشته چطور از پس سربار پیاده‌سازی و افزایش چشمگیر تعداد متخصصان برمی‌آید.

متخصصان مقیم؛ پایانی بر ترافیک سنگین جابجایی وزن‌ها

در پیاده‌سازی قبلی مدل‌های MoE در Ai2، از روش موازی‌سازی داده‌ای کاملاً قطعه‌بندی‌شده (FSDP) استفاده می‌شد؛ ساختاری که در آن وزن‌های متخصصان برای هر دسته داده (بچ کوچک) جمع‌آوری و دوباره خرد می‌شدند. اما وقتی تعداد متخصص‌ها از تعداد کارت‌های گرافیک بیش‌تر می‌شود، جابجایی مداوم این حجم عظیم از وزن‌ها روی بستر ارتباطی کلاستر باعث می‌شود هسته‌های پردازشی دست از کار بکشند و معطل رسیدن داده‌ها بمانند.

پشته OLMo-core 3 با ترکیب یک طراحی مبتنی بر DDP و موازی‌سازی متخصصان، این مشکل را حل کرده است. در این شیوه، متخصصان روی پردازنده‌های گرافیکی اختصاص‌یافته مقیم می‌مانند و این داده‌های توکن هستند که به سمت کارت‌های حاوی متخصصان منتخب حرکت می‌کنند. این چیدمان هوشمندانه، زحمت جمع‌آوری و قطعه‌بندی مجدد تمام وزن‌ها در هر بچ آموزشی را کاملاً از میان برمی‌دارد.

تقسیم سه‌گانه کلاستر برای جا دادن مدل‌های تریلیونی

جا دادن یک مدل با بیش از یک تریلیون پارامتر در حافظه، نیازمند تکه‌تکه کردن هوشمندانه وزن متخصص‌ها، لایه‌های مدل و وضعیت بهینه‌ساز در سراسر کلاستر است:

• موازی‌سازی متخصصان (Expert Parallelism): متخصصان مختلف را بین پردازنده‌های گرافیکی توزیع می‌کند تا هر کارت فقط بخشی از استخر متخصصان را در حافظه خود نگه دارد.
• موازی‌سازی خط لوله‌ای (Pipeline Parallelism): دسته‌هایی از لایه‌های مدل را به مراحل مختلف GPUها اختصاص می‌دهد و حجم وزن‌های نگهداری‌شده روی هر کارت را به حداقل می‌رساند.
• بهینه‌ساز توزیع‌شده (Distributed Optimizer): متغیرها و وضعیت بهینه‌ساز را در سرتاسر GPUها قطعه‌بندی می‌کند تا نیازی به نگهداری یک نسخه کامل روی تک‌تک کارت‌ها نباشد.

مدیریت کامل فرایند مسیریابی روی خود پردازنده گرافیکی

سرعت ارتباطات کلاستر به نحوه بسته‌بندی، انتقال و تحویل توکن‌های مسیریابی‌شده به هر متخصص بستگی دارد. نسخه سوم OLMo-core سه بهینه‌سازی کلیدی در این مسیر اعمال کرده است:

• موازی‌سازی ردیفی متخصصان (Rowwise expert parallelism): ردیف‌های توکن‌های مسیریابی‌شده را مستقیماً درون بافرهای ورودی متخصص می‌نویسد تا بازآرایی و کپی‌های مکرر داده‌ها حذف شود.
• مسیریابی مقیم در پردازنده گرافیکی (GPU-resident routing): متادیتای مسیریابی را روی خود کارت گرافیک نگه می‌دارد؛ در نتیجه پردازنده اصلی (CPU) می‌تواند بدون اینکه منتظر کپی بازگشتی داده از GPU بماند، کار‌های بعدی را سریعاً در صف قرار دهد.
• عملیات ماتریسی گروه‌بندی‌شده (Grouped GEMM): ضرب‌های ماتریسی کوچک و متعدد متخصصان را در قالب عملیات‌های یکپارچه ادغام می‌کند تا از تمام توان پردازشی تراشه به بهینه‌ترین شکل بهره‌برداری شود.

کاهش هزینه‌های پردازش و حافظه با فرمت ممیز شناور MXFP8

فریم‌ورک OLMo-core 3 از فرمت محاسباتی ۸ بیتی MXFP8 پشتیبانی می‌کند؛ قالبی که مقیاس‌گذاری متفاوتی را روی بلوک‌های کوچک اعمال می‌کند. در آزمایشی روی چهار پردازنده گرافیکی NVIDIA B300 با توزیع یکنواخت کار بین متخصصان، استفاده انتخابی از MXFP8 توان عملیاتی آموزش را نسبت به حالت استاندارد BF16 حدود ۲۱ درصد بالا برد و مصرف بیشینه حافظه فعال را از ۱۰۳ گیگابایت به ۹۵ گیگابایت رساند.

بخش عمده این جهش کارایی ناشی از محاسبات فیدفوروارد و ارتباطات سریع‌تر بین متخصصان بود و مکانیزم توجه (Attention) سهم کمتری در آن داشت. البته این سنجش با مسیریابی یکنواخت انجام شده و روترهای یادگیرنده ممکن است بارهای کاری نامتعادل ایجاد کنند؛ بنابراین میزان همگرایی در دوره‌های آموزشی کامل و کیفیت نهایی مدل همچنان نیازمند ارزیابی و اعتبارسنجی مستقل است.

آزمایش غول ۱.۲ تریلیون پارامتری روی ۵۱۲ کارت گرافیک

تیم Ai2 یک مدل عظیم با مجموع ۱.۲ تریلیون پارامتر و ۵۸.۳۶ میلیارد پارامتر فعال به ازای هر توکن را روی ۵۱۲ پردازنده گرافیکی به آزمایش گذاشت. این تست توانست به اوج توان پردازشی ۸۵۸ ترافلاپس در ثانیه به ازای هر GPU (با احتساب محاسبات مفید مدل) دست پیدا کند. در این آزمایش برای سنجش خالص توان فنی سیستم از مسیریابی تصادفی استفاده شد و کیفیت خروجی مدل بررسی نشد.

این پژوهشگران با بهره‌گیری از کتابخانه DeepEP v2 برای تبادل داده بین متخصصان، رکورد ۲.۳۸ تریلیون پارامتر کلی را هم ثبت کردند؛ هرچند این رقم خیره‌کننده حاصل یک آزمایش کوتاه‌مدت برای سنجش سقف ظرفیت بوده و رفتار مدل در آموزش پایدار و پیوسته هنوز گزارش نشده است.

روایت شفاف از چهار شکست و چالش فنی در دل گزارش Ai2

گزارش فنی این مؤسسه چند نتیجه منفی و چالش تجربی را هم با شفافیت مستند کرده که در پیاده‌سازی و بنچمارک مدل‌ها بسیار تأثیرگذارند:

• امتیاز تعادل مسیریابی در حالی بالا می‌رفت که توزیع واقعی بار کاری در عمل نامتعادل‌تر شده بود؛ پدیده‌ای چالش‌برانگیز که تیم تحقیقاتی آن را «جری‌مندرینگ توکن» (Token Gerrymandering - دستکاری متقلبانه مرز توکن‌ها) نامیده است.
• کاهش نرخ یادگیری (Learning Rate) متخصصان برای جبران توکن‌های کمتری که دریافت می‌کردند، هیچ بهبود ملموسی در عملکرد این خانواده از مدل‌ها نشان نداد.
• مدت زمان اجرای دستورات روی GPU حتی با ثابت ماندن ابعاد ماتریس‌ها، با تغییر مقادیر ورودی نوسان می‌کرد؛ بنابراین برای مقایسه دقیق عملکرد، علاوه بر ابعاد یکسان، باید مقادیر ورودی هم کاملاً مطابقت داشته باشند.
• همپوشانی (Overlap) فرآیندهای تبادل داده و محاسبات در استریم‌های مجزای کارت گرافیک، در مواردی حتی مدت‌زمان کل آموزش را طولانی‌تر کرد!

استخرهای بزرگ متخصصان؛ ایده‌آل‌ترین کاربرد OLMo-core 3

پشته OLMo-core 3 دقیقاً برای تیم‌هایی ساخته شده که می‌خواهند بدون افزایش پارامترهای فعال، تعداد متخصص‌ها را بالا ببرند، روش‌های موازی‌سازی پایپ‌لاین و متخصصان را با هم پیوند بزنند، یا کل کدهای آموزش را زیر ذره‌بین بگذارند و سفارشی‌سازی کنند. برای مدل‌های متراکم (Dense) و مدل‌های کوچک MoE که بدون دردسر روی یک کلاستر جا می‌شوند، دلیل چندانی برای مهاجرت از ابزار‌هایی مثل Megatron-Core یا فریم‌ورک‌های مرسوم FSDP وجود ندارد.

کد منبع این پروژه به‌صورت متن‌باز روی گیت‌هاب قرار گرفته و Ai2 یک راهنمای تعاملی تصویری هم منتشر کرده که نشان می‌دهد چطور موازی‌سازی داده‌ای، متخصصان و پایپ‌لاین از یک کارت گرافیک تا ۵۱۲ پردازنده در کنار هم ترکیب می‌شوند.

مؤسسه Ai2 عرضه این زیرساخت‌ها را بخشی از تعهد خود به اکوسیستم مدل‌های باز می‌داند. عمومی کردن این پشته به پژوهشگران امکان می‌دهد از تمام تصمیمات فنی، نتایج بنچمارک‌ها و حتی آزمایش‌های ناموفق پشت صحنه نسل بعدی اولمو سر دربیاورند. این اقدام موانع نرم‌افزاری را برای محققان دانشگاهی و آزمایشگاه‌های کوچک‌تر برمی‌دارد، هرچند که تهیه سخت‌افزارهای غول‌پیکر برای اجرای بزرگ‌ترین پیکربندی‌ها همچنان چالشی جدی به حساب می‌آید.