Olmo-core 3 زیرساخت آموزش مدل‌های MoE را تا مقیاس تریلیونی گسترش می‌دهد

زیرساخت آموزش MoE با انتشار Olmo-core 3 وارد مرحله تازه‌ای می‌شود: AllenAI چارچوب متن‌بازی را معرفی کرده که برای آموزش مدل‌های زبانی بزرگ با معماری «ترکیب متخصصان» بازطراحی شده و در بنچمارک‌ها تا مقیاس بیش از یک تریلیون پارامتر آزمایش شده است. این زیرساخت می‌تواند به پژوهشگران امکان دهد مدل‌های بزرگ‌تری را با استفاده کارآمدتر از GPUها آموزش دهند؛ هرچند برخی از ارقام اعلام‌شده حاصل آزمون عملکرد سامانه‌اند و نباید با آموزش کامل یک مدل در همان مقیاس اشتباه گرفته شوند.

چرا آموزش مدل‌های MoE به زیرساخت ویژه نیاز دارد؟

آموزش مدل‌های هوش مصنوعی بزرگ به منابع محاسباتی فراوان نیاز دارد و همین موضوع هزینه و مصرف انرژی را بالا می‌برد. این محدودیت می‌تواند دسترسی پژوهشگران دانشگاهی و آزمایشگاه‌های کوچک‌تر به توسعه مدل‌های پیشرفته را دشوار کند. مدل‌های MoE یا «ترکیب متخصصان» راهی برای افزایش ظرفیت مدل بدون فعال‌کردن همه اجزای آن برای هر ورودی هستند.

در این معماری، مدل از چندین بخش تخصصی تشکیل می‌شود که به آن‌ها «متخصص» می‌گویند. سامانه برای هر بخش کوچک متن، یعنی هر توکن، تعدادی از این متخصصان را انتخاب می‌کند. بنابراین، مدل ممکن است پارامترهای زیادی داشته باشد، اما در هر مرحله فقط بخشی از آن‌ها محاسبه شود. پارامترها مقادیری‌اند که مدل هنگام آموزش یاد می‌گیرد و برای تولید پاسخ به کار می‌برد.

فعال‌کردن بخشی از مدل به‌تنهایی همه دشواری‌ها را برطرف نمی‌کند. کل مدل و وضعیت‌های لازم برای آموزش آن همچنان باید در حافظه GPUها جا بگیرند. افزون بر این، ورودی‌ها باید به متخصص مناسب در میان چند GPU یا چند سرور برسند. جابه‌جایی داده و هماهنگ‌کردن محاسبات می‌تواند زمان‌بر باشد و بخشی از صرفه‌جویی معماری MoE را از بین ببرد.

به گفته AllenAI، Olmo-core 3 برای کاهش این فاصله طراحی شده است. در یکی از بنچمارک‌ها، تعداد متخصصان از ۸ به ۱۲۸ افزایش یافت، اما برای هر توکن همچنان فقط چهار متخصص انتخاب شدند. شمار پارامترهای فعال برای هر توکن نزدیک به ۳٫۲ میلیارد باقی ماند؛ در همین حال، ظرفیت کل پارامترها از ۴٫۶ میلیارد به ۴۷ میلیارد رسید و سرعت آموزش کمتر از ۵ درصد کاهش پیدا کرد.

زیرساخت آموزش MoE و تغییر شیوه توزیع محاسبات

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

یکی از تغییرهای مهم، جایگزینی روش پیشین مبتنی بر FSDP با سامانه‌ای مبتنی بر DDP است. در پیاده‌سازی قبلی، وزن‌های مدل برای دسته‌های کوچک داده گردآوری و دوباره میان پردازنده‌ها توزیع می‌شدند. نسخه تازه، متخصصان را روی GPUها نگه می‌دارد و داده مرتبط را به آن‌ها می‌فرستد؛ به این ترتیب، نیاز به گردآوری مکرر وزن‌ها کاهش می‌یابد.

این رویکرد با چند روش مکمل همراه است. «موازی‌سازی متخصصان» متخصصان را میان GPUها پخش می‌کند تا هر GPU فقط بخشی از مجموعه را نگه دارد. «موازی‌سازی خط لوله» لایه‌های مدل را میان گروه‌های GPU تقسیم می‌کند. همچنین، بهینه‌ساز توزیع‌شده وضعیت‌های لازم برای به‌روزرسانی مدل را میان GPUها پخش می‌کند، به‌جای آن‌که هرکدام نسخه کاملی از این داده‌ها را ذخیره کنند.

AllenAI همچنین برای کاهش هزینه مسیریابی داده و اجرای محاسبات، از روش‌هایی مانند موازی‌سازی سطری متخصصان استفاده کرده است. این روش داده مسیریابی‌شده را مستقیم در بافر ورودی متخصصان قرار می‌دهد و کار اضافی برای جابه‌جایی آن را کم می‌کند. مسیریابی مقیم GPU نیز اطلاعات مسیریابی را روی GPU نگه می‌دارد تا پردازنده مرکزی برای دریافت دوباره آن منتظر نماند. در کنار این موارد، grouped GEMM چند محاسبه کوچک متخصصان را با هم ترکیب می‌کند تا GPUها آن‌ها را کارآمدتر اجرا کنند.

نتایج بنچمارک و نقش قالب MXFP8

در آزمونی مقدماتی روی هشت NVIDIA B300، یک مدل MoE با ۴۷ میلیارد پارامتر در Olmo-core 3 به سرعت ۵۲ هزار توکن در ثانیه به‌ازای هر GPU رسید. پیاده‌سازی پیشین AllenAI در همان مقایسه ۱۹ هزار و ۴۰۰ توکن در ثانیه به‌ازای هر GPU پردازش کرده بود. طبق اعلام این گروه، نتیجه نسخه جدید حدود ۲٫۷ برابر سرعت بیشتر است. این مقایسه به همان آزمون مقدماتی مربوط می‌شود و به‌تنهایی درباره کیفیت یا توانایی نهایی یک مدل چیزی نمی‌گوید.

یکی دیگر از بهینه‌سازی‌ها پشتیبانی از MXFP8 است؛ قالبی با دقت پایین‌تر که بعضی مقادیر را با بیت‌های کمتری نمایش می‌دهد. کاهش تعداد بیت‌ها می‌تواند محاسبه و جابه‌جایی داده میان GPUها را سبک‌تر کند، اما تبدیل میان قالب‌های عددی نیز هزینه دارد. بنابراین، استفاده از دقت پایین‌تر زمانی سودمند است که صرفه‌جویی آن از هزینه تبدیل بیشتر باشد.

در یک آزمون کنترل‌شده روی چهار NVIDIA B300، فعال‌کردن MXFP8 در بخش‌هایی که بیشترین سود را داشتند، سرعت آموزش را در مقایسه با قالب BF16 حدود ۲۱ درصد افزایش داد. حافظه فعال اوج نیز از ۱۰۳ گیبی‌بایت به ۹۵ گیبی‌بایت کاهش یافت. به گفته AllenAI، بخش عمده بهبود به محاسبات پیش‌خور و انتقال داده میان متخصصان مربوط بود، نه فقط به سازوکار توجه.

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

مقیاس تریلیونی؛ چه چیزی اندازه‌گیری شده است؟

AllenAI می‌گوید Olmo-core 3 را روی پیکربندی‌های گوناگون NVIDIA B300 آزمایش کرده است. یکی از آن‌ها مدلی با ۱٫۲ تریلیون پارامتر کل بود که روی ۵۱۲ GPU اجرا شد و ۵۸٫۳۶ میلیارد پارامتر را برای هر توکن فعال می‌کرد. بیشترین سرعت ثبت‌شده در این آزمون ۸۵۸ ترافلاپس به‌ازای هر GPU بود؛ معیاری برای سنجش حجم محاسبات مفید مدل در هر ثانیه.

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

گزارش فنی پروژه علاوه بر سنجش سرعت، به تجربه‌هایی درباره چگونگی آموزش MoE می‌پردازد. برای نمونه، پژوهشگران پدیده‌ای را «دست‌کاری توکن» نامیده‌اند: امتیازی که قرار بود توزیع متوازن ورودی‌ها میان متخصصان را تشویق کند، گاهی بهتر می‌شد، حتی وقتی بار واقعی متخصصان نامتوازن‌تر بود. آن‌ها همچنین گزارش کرده‌اند که کاهش نرخ یادگیری متخصصان، با این استدلال که هرکدام توکن‌های کمتری پردازش می‌کنند، در خانواده مدلی که آزمایش کردند به بهبود منجر نشد.

برخی نتایج دیگر بر ضرورت دقت در مقایسه عملکرد تأکید دارند. زمان محاسبات GPU ممکن است با تغییر مقادیر ورودی فرق کند، حتی اگر ابعاد ماتریس‌ها یکسان باشد؛ بنابراین، مقایسه‌ها باید هم شکل و هم مقادیر ورودی را در نظر بگیرند. همچنین، هم‌پوشانی ارتباطات و محاسبات در چند جریان GPU در همه آزمون‌ها سرعت را بالا نبرد. این موارد نشان می‌دهد که افزایش مقیاس صرفاً با افزودن GPU حل نمی‌شود و جزئیات اجرای سامانه اهمیت دارد.

برنامه AllenAI برای نسل بعدی Olmo

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

زیرساخت تازه کاملاً متن‌باز معرفی شده است. پژوهشگران و توسعه‌دهندگان می‌توانند از آن برای آموزش مدل‌های MoE خود، سازگارکردن چارچوب با سخت‌افزارهای متفاوت و آزمایش روش‌های مسیریابی و موازی‌سازی استفاده کنند. بازبودن ابزار آموزش، در کنار انتشار وزن‌های مدل، می‌تواند امکان بررسی چگونگی ساخته‌شدن مدل و آزمایش تصمیم‌های فنی پشت آن را فراهم کند؛ هرچند کارایی نهایی به سخت‌افزار، داده و شیوه استفاده نیز وابسته است.

این خبر برای چه کسانی اهمیت دارد

Olmo-core 3 بیش از همه برای گروه‌های دانشگاهی، آزمایشگاه‌های کوچک و توسعه‌دهندگانی اهمیت دارد که می‌خواهند مدل‌های MoE را آموزش دهند یا درباره زیرساخت آن‌ها پژوهش کنند. نتایج اعلام‌شده نشان می‌دهد که افزایش شمار متخصصان و ظرفیت کل مدل الزاماً به افت شدید سرعت نمی‌انجامد، اما آزمون‌های مقیاس تریلیونی را باید به‌عنوان سنجش ظرفیت سامانه خواند، نه اثبات آموزش موفق یک مدل کامل در آن اندازه.

برای فعالان زیرساخت و پژوهشگران یادگیری ماشین نیز گزارش فنی می‌تواند ارزشمند باشد، چون علاوه بر ارقام عملکرد، نمونه‌هایی از بهینه‌سازی‌های ناموفق یا محدودیت‌های اندازه‌گیری را شرح می‌دهد. در مجموع، اهمیت این انتشار در ارائه ابزار و داده‌های فنی باز برای کار روی آموزش MoE است؛ نه در معرفی یک مدل زبانی تازه برای استفاده عمومی.