Olmo-core 3 توان عملیاتی مدل‌های تریلیونی را با معماری MoE افزایش داد

اگر قصد دارید مدلی با تریلیون‌ها پارامتر آموزش دهید، احتمالاً با این واقعیت تلخ روبرو شده‌اید که سخت‌افزار شما بیشتر از آنکه محاسبه کند، مشغول جابه‌جایی داده‌هاست. مؤسسه Allen Institute for AI (Ai2) در ۱ اکتبر ۲۰۲۶ با معرفی Olmo-core 3، راهکاری برای حذف این «مالیات ارتباطی» در مدل‌های عظیم ارائه کرد. این پشته آموزشی بازطراحی شده، امکان بنچمارک کردن یک مدل ۱.۲ تریلیون پارامتری را با توان عملیاتی بالا و با استفاده از یک چارچوب متن‌باز فراهم می‌کند.

طبق اعلام این مؤسسه، چالش اصلی در معماری ترکیب خبره‌ها (Mixture of Experts یا MoE) — که شبیه به شرکتی است که هر درخواست را به متخصص‌ترین کارمندش می‌سپارد تا کل کارکنان درگیر نشوند — هزینه بالای هماهنگی بین پردازنده‌هاست. در مدل‌های MoE استاندارد، مسیریابی داده‌ها به سمت «خبره» درست در یک خوشه GPU، اغلب همان بازدهی‌ای را که مدل‌های پراکنده (Sparse) ایجاد می‌کنند، از بین می‌برد. Olmo-core 3 دقیقاً برای پر کردن این شکاف طراحی شده است تا با رشد مدل، زمان صرف‌شده برای جابه‌جایی داده‌ها از زمان محاسبات پیشی نگیرد.

زمینه و تکامل

کاوش‌های Ai2 در زمینه مدل‌های پراکنده با OlmoE آغاز شد که از یک معماری MoE با ۶۴ خبره مسیریابی شده استفاده می‌کرد. در مقابل، Olmo 3 از یک معماری متراکم (Dense) بهره می‌برد که در آن تقریباً تمام مدل برای هر توکن فعال می‌ماند. Olmo-core 3 اکنون این چارچوب را با ارائه یک سیستم آموزشی که به‌طور خاص برای مدل‌های MoE بسیار بزرگ‌تر طراحی شده است، گسترش می‌دهد.

در حالی که Megatron-Core شرکت NVIDIA پیش از این یک گزینه تثبیت‌شده برای آموزش MoEهای بزرگ بود، Ai2 توصیف می‌کند که Olmo-core 3 یک پشته آموزشی MoE کاملاً یکپارچه را مستقیماً به چارچوب پشتیبان Olmo می‌آورد. این انتشار با یک دموی تعاملی، کد باز و یک گزارش فنی با عنوان «Supercharging Olmo-core for Efficient and Scalable MoE Training» پشتیبانی می‌شود. این گزارش توسط پژوهشگران Ai2 و دانشگاه واشینگتن نوشته شده و Tianhua Tao نویسنده اصلی و توسعه‌دهنده کلیدی آن است.

از تکه‌تکه سازی تا خبره‌های مقیم

تیم توسعه‌دهنده به رهبری Tianhua Tao، رویکرد بنیادین تعامل مدل با سخت‌افزار را تغییر داده است. نسخه‌های قبلی از موازات داده‌های کاملاً تکه‌تکه شده (FSDP) استفاده می‌کردند که در آن وزن‌های مدل برای هر دسته کوچک از داده‌ها باید جمع‌آوری و دوباره توزیع می‌شدند.

در مقابل، Olmo-core 3 از موازات داده‌های توزیع‌شده (DDP) استفاده می‌کند. در این روش، خبره‌ها به‌صورت ثابت (Resident) روی GPUها مستقر می‌شوند و این داده‌ها هستند که به سمت آن‌ها مسیریابی می‌شوند. در یک تست با ۸ پردازنده NVIDIA B300، یک مدل MoE با ۴۷ میلیارد پارامتر توانست ۵۲,۰۰۰ توکن در ثانیه به ازای هر GPU پردازش کند؛ این رقم تقریباً ۲.۷ برابر بیشتر از پیاده‌سازی قبلی است که تنها ۱۹,۴۰۰ توکن را مدیریت می‌کرد. این تلاش برای بهینه‌سازی زیرساخت، در راستای موج جدیدی از متدهای آموزشی است که با استفاده از دستورالعمل‌های پیش‌آموزش بهینه، هزینه‌های محاسباتی مدل‌های باز را تا ۱۰ برابر کاهش داده‌اند.

مقیاس‌بندی استخر خبره‌ها

یکی از خیره‌کننده‌ترین بنچمارک‌ها، گسترش تعداد خبره‌ها از ۸ به ۱۲۸ مورد بود، در حالی که پارامترهای فعال برای هر توکن روی ۳.۲ میلیارد ثابت ماند. با وجود اینکه ظرفیت کل پارامترها از ۴.۶ میلیارد به ۴۷ میلیارد رسید، توان عملیاتی (Throughput) کمتر از ۵٪ کاهش یافت. این موضوع ثابت می‌کند که سیستم می‌تواند ظرفیت را بدون افزایش خطی هزینه‌های محاسباتی بالا ببرد.

جزئیات موازات و حافظه

برای استخراج حداکثری قدرت سخت‌افزار، Olmo-core 3 از سه تکنیک اصلی برای تقسیم مدل و وضعیت آموزش در سخت‌افزار استفاده می‌کند:

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

به گزارش این تیم، بهینه‌سازی‌های فنی دیگری نیز اعمال شده است:

• موازات خبره ردیفی (Rowwise Expert Parallelism): استفاده از طراحی مبتنی بر NVSHMEM برای نوشتن مستقیم توکن‌ها در بافرهای ورودی خبره، که این فرآیند را بدون نیاز به همگام‌سازی (Synchronization-free) می‌کند.
• مسیریابی مستقر در GPU: متادیتای مسیریابی روی GPU می‌ماند تا CPU بتواند بدون انتظار برای کپی داده‌ها، صف کاری را مدیریت کند.
• Grouped GEMM: ترکیب چندین محاسبه کوچک خبره در اجراهای بزرگ‌تر و بهینه‌تر GPU.
• پشتیبانی از MXFP8: استفاده از این فرمت با دقت پایین‌تر، توان عملیاتی آموزش را در مقایسه با BF16 روی چهار GPU NVIDIA B300 حدود ۲۱٪ افزایش داد.
• چک‌پوینت‌های مستقل از توپولوژی: این سیستم تنسورهای جهانی FP32 را مستقل از چیدمان موازات ثبت می‌کند، که اجازه می‌دهد یک اجرای متوقف شده روی توپولوژی سخت‌افزاری متفاوتی از سر گرفته شود.

مدیریت حافظه نیز بهبود یافت. حافظه فعال پیک (Peak active memory) هنگام استفاده از MXFP8 از ۱۰۳ گیگابایت به ۹۵ گیگابایت کاهش یافت که بیشتر این دستاورد مربوط به محاسبات Feed-forward و جابه‌جایی داده‌ها بود. علاوه بر این، بازمحاسبه فعال‌سازها (Activation Recompute) تقریباً دو-سوم از حافظه فعال‌ساز را حذف کرد و حافظه پیک را حدود یک چهارتم کاهش داد، هرچند که این کار توان عملیاتی را حدود یک پنجم کاهش داد.

فشار به مرز تریلیون پارامتر

طبق گزارش فنی، تیم سیستم را تا مرزهای افراطی پیش برد. آن‌ها یک مدل ۱.۲ تریلیون پارامتری را روی ۵۱۲ پردازنده تست کردند که در آن ۵۸.۳۶ میلیارد پارامتر به ازای هر توکن فعال بود و به توان عملیاتی ۸۵۸ TFLOP/s به ازای هر GPU رسیدند.

Ai2 همچنین از DeepEP v2 به عنوان یک بک‌اند ارتباطی جایگزین استفاده کرد. این ابزار به آن‌ها اجازه داد تا در یک تست ظرفیت کوتاه، به ۲.۳۸ تریلیون پارامتر کل دست یابند. تیم اشاره کرد که در این تست‌های سطح بالا، برای اندازه‌گیری عملکرد سیستم از مسیریابی تصادفی استفاده شده است و نه برای سنجش کیفیت واقعی یک مدل آموزش‌دیده. گزارش نقاط عملیاتی را از یک مدل ۱۲.۹ میلیارد پارامتری روی ۱۶ GPU تا پیکربندی ۱.۲ تریلیون پارامتری فهرست کرده است.

شکست‌های مستند شده و یافته‌ها

این گزارش در مورد مواردی که نتیجه نداده‌اند نیز بسیار شفاف است. پژوهشگران پدیده‌ای به نام «تجزیه توکن‌ها» (Token Gerrymandering) را شناسایی کردند؛ جایی که نمرات مسیریابی که برای تشویق تعادل طراحی شده بودند، بهبود می‌یافتند اما در واقعیت، حجم کاری بین خبره‌ها نامتوازن باقی می‌ماند.

سایر یافته‌های مستند شده عبارتند از:

• نرخ‌های یادگیری: کاهش نرخ یادگیری برای خبره‌هایی که توکن‌های کمتری پردازش می‌کردند، در خانواده مدل‌های تست شده منجر به بهبود نتایج نشد.
• واریانس GPU: محاسبات بر اساس مقادیر در حال پردازش، زمان‌های متفاوتی می‌برد، حتی زمانی که ابعاد ماتریس‌ها یکسان بود.
• هم‌پوشانی جریان (Stream Overlap): هم‌پوشانی ارتباطات و محاسبات روی جریان‌های مجزای GPU گاهی اوقات باعث کند شدن اجرای کلی (End-to-end) می‌شد.
• رویکردهای رد شده: تیم روش انتقال فعال‌سازها به CPU (به دلیل محدودیت‌های لینک میزبان) و دو طرح هم‌پوشانی که با محاسبات خبره برای منابع GPU رقابت می‌کردند را تست کرد اما نپذیرفت.

این زیرساخت اکنون پایه و اساس نسل بعدی Olmo است که قرار است با معماری MoE، بزرگ‌ترین مجموعه داده و طولانی‌ترین پنجره متنی ساخته شود. هدف Ai2 این است که این نسخه توانمندترین مدل آن‌ها باشد.

برای توسعه‌دهندگان، Olmo-core 3 کاملاً تحت لایسنس Apache-2.0 باز است. این ابزار به صورت بلوک‌های سازنده PyTorch از طریق GitHub یا PyPI با نام ai2-olmo-core در دسترس است.

این تغییر به سمت پشته‌های آموزشی متن‌باز تریلیون پارامتری، پیش‌فرض‌های این حوزه را درباره اینکه چه کسی می‌تواند مدل‌های مقیاس مرزی (Frontier-scale) بسازد، تغییر می‌دهد. Ai2 با حذف «مالیات ارتباطی» MoEها، مانع پژوهشگران را برای آزمایش پراکندگی‌های شدید بدون نیاز به لایه‌های ارکستراسیون بسته و اختصاصی کاهش می‌دهد.

پژوهشگران اکنون باید انتشار مدل Olmo MoE را زیر نظر بگیرند تا ببینند آیا این افزایش توان عملیاتی به هوش مدل و قابلیت‌های استدلال برتر تبدیل می‌شود یا خیر.

گام بعدی شما

• اگر توسعه‌دهنده مدل‌های بزرگ هستید، کتابخانه ai2-olmo-core را برای جایگزینی FSDP با DDP در مدل‌های MoE خود تست کنید.
• گزارش فنی «Supercharging Olmo-core» را برای درک جزئیات پیاده‌سازی MXFP8 و مدیریت حافظه مطالعه کنید.
• منتظر انتشار نسخه MoE مدل Olmo باشید تا ببینید آیا این افزایش سرعت به بهبود استدلال مدل منجر شده است یا خیر.

اما تأثیر این بهینه‌سازی‌ها بر هزینه‌های استنتاج در مقیاس تجاری حتی پیچیده‌تر است — به تحلیل ما درباره‌ی کاهش هزینه‌های Inference در مدل‌های پراکنده مراجعه کنید.