معماری "Mixture-of-Experts"؛ موتور قدرتمندترین LLMهای تولیدی — و عملکردی فراتر از تصور رایج
هنگامی که OpenAI مدل GPT-4 را عرضه کرد، از اعلام تعداد پارامترها خودداری کرد. ماهها بعد، اسناد لو رفته و Benchmarkهای تأییدکننده نشان داد که این مدل از معماری Mixture-of-Experts (MoE) با حدود ۱٫۸ تریلیون پارامتر کل در هشت زیرشبکه تخصصی استفاده میکند — اما در هر پاس رو به جلو تنها حدود ۲۲۰ میلیارد پارامتر فعال میشود. همین یک تصمیم طراحی، هم سقف توانایی مدل و هم اقتصاد Inference آن را به گونهای توضیح میدهد که یک تعداد پارامتر ساده هرگز نمیتوانست.
اکنون معماری MoE برای مدلهای مرزی غالب است. Gemini 1.5 گوگل از MoE استفاده میکند. مدلهای منبعباز Mixtral 8x7B و 8x22B شرکت Mistral AI، MoE را برای میزبانان خودکار در دسترس قرار دادند. تحقیقات داخلی متا روی MoE برای جانشینهای Llama نیز مستند شده است. درک این که MoE واقعاً چگونه کار میکند — و در کجا واقعاً کمک میکند در مقابل جایی که فقط اسلایدها را جذابتر نشان میدهد — اگر تصمیم دارید کدام مدلها را مستقر کنید یا انتشارات جدید را ارزیابی کنید، اهمیت دارد.
ایده اصلی: محاسبات شرطی
یک Transformer متراکم استاندارد مانند Llama 2 70B تمام ۷۰ میلیارد پارامتر خود را برای هر Token پردازش شده فعال میکند. این کار از نظر محاسباتی پرهزینه ولی قابل پیشبینی است. معماری MoE لایههای Feedforward (لایههایی که بخش عمده پارامترهای یک Transformer را تشکیل میدهند) را با چندین شبکه "متخصص" موازی به همراه یک مسیریاب سبک جایگزین میکند. برای هر Token، مسیریاب بهترین k متخصص را انتخاب میکند — معمولاً ۲ مورد از ۸ یا ۱۶ — و تنها آن متخصصان آن Token را پردازش میکنند. نتایج وزندهی شده و ترکیب میشوند.
نتیجه عملی: مدل Mixtral 8x7B دارای حدود ۴۷ میلیارد پارامتر کل است، اما هر Token تنها به حدود ۱۳ میلیارد از آنها دست میزند. شما بیشتر ظرفیت نمایشی یک مدل متراکم ۴۷ میلیاردی را با هزینه Inference نزدیک به ۱۳ میلیارد دریافت میکنید. توان عملیاتی در مقایسه با یک مدل متراکم معادل روی سختافزار یکسان برای خروجی با کیفیت یکسان، تقریباً دو برابر میشود.
مسیریاب واقعاً چه چیزی یاد میگیرد
مسیریاب یک لایه خطی کوچک است که توزیع احتمالی روی همه متخصصان موجود تولید میکند. این لایه به صورت سرتاسری با بقیه مدل با استفاده از gradient descent استاندارد آموزش داده میشود — هیچ پیشآموزش یا برچسبگذاری دستی جداگانهای برای این که کدام متخصص باید چه محتوایی را مدیریت کند وجود ندارد. آنچه پدیدار میشود تقریباً تخصص حوزهای است: تحلیل الگوهای مسیریابی Mixtral نشان میدهد که متخصصان ترجیحات نرمی برای نحو کد، استدلال زبان طبیعی، یادآوری واقعیات و غیره پیدا میکنند. اما این تخصص دقیق نیست و همیشه با شهود انسانی درباره موضوع هماهنگ نیست.
یک مشکل مهندسی پایدار، توازن بار است. بدون مداخله، مسیریاب تمایل دارد روی مجموعه کوچکی از متخصصان "محبوب" متمرکز شود و دیگران را گرسنه بگذارد و ظرفیت را هدر دهد. راه حل استاندارد یک تلفات کمکی توازن بار است که در طول آموزش اضافه میشود و استفاده نابرابر از متخصصان را جریمه میکند. تنظیم قدرت این تلفات یک Hyperparameter است که هم کیفیت مدل و هم کارایی سختافزار را به طور معناداری تحت تأثیر قرار میدهد — بیش از حد کم باعث فروپاشی متخصصان میشود؛ بیش از حد زیاد، مسیریاب نمیتواند تخصص معنادار یاد بگیرد.
گلوگاه حافظهای که بازاریابی نادیده میگیرد
اینجا جایی است که MoE برای مستقرکنندگان پیچیده میشود. همه پارامترها باید در حافظه قرار گیرند حتی اگر فقط بخشی از آنها به ازای هر Token فعال شود. مدل Mixtral 8x22B — با حدود ۱۴۱ میلیارد پارامتر کل — تقریباً به ۲۸۰ گیگابایت VRAM GPU در دقت BF16 نیاز دارد قبل از این که KV cache را حساب کنید. این یعنی حداقل چهار GPU H100 با ۸۰ گیگابایت فقط برای نگه داشتن وزنها، حتی اگر توان عملیاتی Inference مشابه یک مدل متراکم بسیار کوچکتر باشد.
این یک شکاف زیرساختی ایجاد میکند. در یک مرکز داده که میتوانید یک گره ۴-GPU را به هر نسخه از مدل اختصاص دهید، MoE واقعاً به ازای هر Token ارزانتر است. در استقرارهایی که میخواهید چندین مدل را روی سختافزار مشترک مستقر کنید، ردپای حافظه MoE آن را گران میکند. همچنین به همین دلیل است که quantization برای مدلهای MoE مهمتر است: رساندن Mixtral 8x7B به دقت ۴ بیتی (تقریباً ۲۵ گیگابایت) آن را برای اجرا روی یک ایستگاه کاری مصرفگرا یا یک سرور دو-GPU عملی میکند.
موازیسازی متخصصان به عنوان اهرم مقیاسدهی
برای آموزش مدلهای MoE بسیار بزرگ، تکنیکی به نام expert parallelism متخصصان مختلف را در GPUهای فیزیکی مختلف توزیع میکند. هنگامی که یک Token به متخصص شماره ۵ هدایت میشود، محاسبه روی GPUای که وزنهای متخصص ۵ را نگه میدارد انجام میشود و نتیجه بازگردانده میشود. این کار ارتباط all-reduce را به انتقالات محلی نقطهبهنقطه تبدیل میکند و آموزش را در مقیاسهایی ممکن میسازد که در غیر این صورت به حافظه per-GPU بیش از حد نیاز داشت.
مقاله Switch Transformer گوگل در سال ۲۰۲۱ این را با ۱٫۶ تریلیون پارامتر نشان داد — اولین مدل تریلیون پارامتری مستند عمومی. یافته کلیدی: یک MoE با ۶۴ متخصص با بودجه محاسباتی مشابه یک مدل متراکم T5-XXL به ۴ برابر سرعت در زمان آموزش دست یافت در حالی که کیفیت را در Benchmarkهای استاندارد مطابقت داد یا فراتر برد. این مقاله همچنین حالتهای شکست را مستند کرد: ناپایداری آموزش در تعداد متخصصان بالا، مشکل فروپاشی توازن بار، و سربار ارتباطات در تنظیمات چندگره.
جایی که MoE واقعاً از مدلهای متراکم پایینتر است
یادگیری few-shot در وظایف بسیار خاص دامنه یکی از مناطقی است که مدلهای MoE ممکن است از مدلهای متراکم هماندازه عملکرد ضعیفتری داشته باشند. از آنجا که مسیریاب Tokenها را به صورت احتمالی تخصیص میدهد و Tokenهای مختلف در یک Prompt ممکن است به متخصصان مختلف بروند، "حافظه" مدل از زمینه اولیه میتواند در بین متخصصان تکهتکه شود و به صورتهایی که انسجام را در اسناد طولانی و تخصصی آسیب میزند. گزارشهای حکایتی از استقرارهای سازمانی Mixtral نشان میدهد که مدلهای متراکم با هزینه Inference معادل گاهی نتایج بهتری بر روی متون حقوقی یا پزشکی که ثبات اصطلاحات دقیق مهم است تولید میکنند.
اندازه Batch نیز مهم است. مزیت توان عملیاتی معماری MoE بیشتر در اندازههای Batch بزرگ که همه متخصصان استفاده تقریباً یکنواخت دارند، آشکار است. در اندازه Batch 1 — یک کاربر واحد که یک query بلادرنگ میدهد — شما دو متخصص را فعال میکنید و شش متخصص دیگر بیکار میمانند. تأخیر به ازای هر Token میتواند به دلیل سربار مسیریابی بدتر از یک مدل متراکم با تعداد پارامتر فعال معادل باشد. به همین دلیل است که استقرارهای تولیدی درخواستها را به شدت Batch میکنند و API endpoints استریم پروفایل تأخیر متفاوتی نسبت به endpoints Inference Batch دارند.
تصمیمات عملی برای تیمهایی که مدلهای MoE را ارزیابی میکنند
اگر یک مدل متراکم ۷۰B را با یک مدل MoE مانند Mixtral 8x22B برای استقرار مقایسه میکنید، مقایسه درست تعداد پارامترها نیست — بلکه ردپای حافظه در برابر کیفیت روی بار کاری خاص شماست. هر دو را روی توزیع وظایف واقعی خود اجرا کنید. Mixtral 8x22B به طور پایدار Llama 2 70B را در Benchmarkهای استدلال شکست میدهد اما شکاف در وظایف narrow retrieval-augmented generation که مجموعه داده همگن است به طور قابل توجهی کاهش مییابد.
برای fine-tuning، مدلهای MoE یک چالش خاص ارائه میدهند: LoRA fine-tuning که فقط روی لایههای متراکم اعمال میشود به وزنهای متخصص دست نمیزند که شامل اکثر دانش تخصصی مدل است. Fine-tuning کامل مدلهای MoE حافظهبر است. انواع خاص MoE از LoRA که آداپتورها را به لایههای Feedforward متخصص اعمال میکنند وجود دارند اما هنوز ابزار استاندارد نیستند — قبل از تعهد، بررسی کنید که چارچوب fine-tuning شما از آنها پشتیبانی میکند یا خیر.
خود وزنهای مسیریاب میتوانند در طول fine-tuning منجمد شوند تا الگوهای تخصص که در پیشآموزش یاد گرفته شدهاند حفظ شوند. این زمانی خوب کار میکند که fine-tuning برای وظیفهای انجام شود که در توزیع آموزش اصلی به خوبی بازنمایی شده است. هنگام تطبیق با یک حوزه واقعاً جدید، باز کردن مسیریاب و پذیرش زمان بیشتر fine-tuning ارزش آن را دارد.
آینده چه خواهد شد
جهتهای تحقیقاتی که به طور فعال بررسی میشوند شامل MoE پراکنده با بیش از دو متخصص فعال به ازای هر Token (مبادله محاسبه با کیفیت)، مسیریابی سلسلهمراتبی که یک مسیریاب درشت ابتدا "خانوادههای" متخصص را انتخاب میکند سپس یک مسیریاب دقیق متخصصان خاص را انتخاب میکند، و معماریهای mixture-of-depths که Tokenها را به لایههای مختلف به جای متخصصان مختلف در یک لایه هدایت میکند، هستند. مقاله گوگل DeepMind در سال ۲۰۲۴ درباره mixture-of-depths نشان داد که هر Token نیازی به عبور از هر لایه Transformer ندارد و امکان دستاوردهای محاسبات شرطی بیشتر را فراهم میکند.
درس معماری از MoE ثابت است: قوانین مقیاسدهی به محاسبات شرطی پاداش میدهند. صرف تمام محاسبات روی هر Token برای هر کار هدر دادن است. مدلهایی که در دو سال آینده اهمیت خواهند داشت به طور فزاینده سیستمهای ترکیبی خواهند بود که کار را هوشمندانه هدایت میکنند — چه هدایت به متخصصان درون یک مدل، چه هدایت به مدلهای مختلف از طریق orchestration، یا هدایت به ابزارهای خارجی. MoE اولین نمایش در مقیاس تولید است که این اصل در سطح وزنها کار میکند.