SkewAdam حافظه آموزش Mixture-of-Experts را بیش از ۶۰٪ کاهش میدهد و جهش قابلتوجهی در دقت ایجاد میکند که به توسعهدهندگان اجازه میدهد یک مدل MoE با ۶.۷۸ میلیارد پارامتر را روی یک GPU واحد ۴۰ گیگابایتی جای دهند.
چرا آموزش MoE با سد حافظه مواجه میشود
معماریهای Mixture-of-Experts هر ورودی را از طریق زیرمجموعه کوچکی از زیرشبکههای «expert» هدایت میکنند، در حالی که یک ستون فقرات متراکم (dense backbone) را حفظ مینمایند. مزیت این کار، داشتن مدلی است که میتواند بدون افزایش متناسب در محاسبات، تا میلیاردها پارامتر مقیاسپذیر شود. در عمل، بهینهساز (optimizer) — بهویژه نسخه AdamW که اکثر تیمها از آن استفاده میکنند — بخش عمدهای از RAM گرافیکی را اشغال میکند. برای یک مدل با ۶.۷۸ میلیارد پارامتر، تنها تنسورهای مومنتوم (momentum) و واریانس (variance) بهینهساز حدود ۵۰ گیگابایت فضا میخواهند. با اضافه شدن اکتیویشنها (activations) و وزنهای مدل، اوج مصرف حافظه به ۸۱.۴ گیگابایت میرسد که استفاده از تنظیمات چند-GPU یا برنامههای آموزشی کندتر را اجباری میکند.
تفاوت SkewAdam در چیست
SkewAdam قانون یادگیری جدیدی ابداع نمیکند؛ بلکه محل قرارگیری لحظات (moments) Adam را بازآرایی میکند:
- ستون فقرات متراکم (dense backbone)، مومنتوم با دقت کامل (full-precision) را حفظ میکند تا بهروزرسانیهای نرم مورد نیاز لایههای متراکم حفظ شود.
- بانک متخصصان (expert bank)، یک تقریب فاکتورگیریشده از واریانس را ذخیره میکند که باعث کاهش دادههای نگهداری شده برای هر متخصص میشود.
- روتر (router)، که تصمیم میگیرد کدام متخصصها فعال شوند، تخمین دقیق لحظه دوم (second-moment) را حفظ میکند.
با در نظر گرفتن این سه جزء به عنوان «سطوح» (tiers) مجزا، این بهینهساز دقت اضافی و مازاد را در جاهایی که اهمیت کمتری دارد حذف کرده و آن را در جاهایی که بیشترین اهمیت را دارد حفظ میکند.
تأثیر قابل اندازهگیری
اجرای همان مدل MoE با ۶.۷۸ میلیارد پارامتر با استفاده از SkewAdam نتایج زیر را به همراه دارد:
- اوج حافظه GPU: ۳۱.۳ گیگابایت (کاهش یافته از ۸۱.۴ گیگابایت)
- میزان اشغال فضای وضعیت بهینهساز: ۱.۲۹ گیگابایت (کاهش یافته از ۵۰ گیگابایت)
این وضعیت کاهشیافته بهراحتی در یک شتابدهنده ۴۰ گیگابایتی واحد جای میگیرد.
بهبود دقت، نه فقط صرفهجویی
کاهش حافظه اغلب به کیفیت مدل آسیب میزند، اما SkewAdam میزان پرپلکسیتی (perplexity) — یک معیار استاندارد در مدلهای زبانی — را از ۱۲۶.۸ (در AdamW) به ۱۰۸.۴ بهبود میبخشد. همچنین این روش در همان وظیفه، از Muon (۱۲۰.۲) و Lion (۳۹۳.۷) عملکرد بهتری دارد. نویسندگان میگویند این بهبود ناشی از حفظ مومنتوم در هر سه سطح است؛ روشهایی که مومنتوم را بهطور کامل حذف میکنند، مانند Adafactor، در این زمینه عقب میمانند.
پرسشهای بیپاسخ
نتایج SkewAdam روی یک مدل ۶.۷۸ میلیارد پارامتری نشان داده شده است. هنوز مشخص نیست که آیا همان نسبتهای وضعیت حافظه برای مدلهایی با یک مرتبه بزرگی (order of magnitude) بزرگتر یا برای وظایفی فراتر از مدلسازی زبان نیز صادق است یا خیر.
آنچه باید زیر نظر داشت
- بنچمارکها روی پیکربندیهای بزرگتر MoE (دهها میلیارد پارامتر).
- پذیرش توسط فریمورکهای متنباز و گنجانده شدن در اسکریپتهای آموزشی محبوب.
- بازخوردهای جامعه علمی درباره موازنههای (trade-offs) پنهان، مانند سرعت همگرایی (convergence speed) یا پایداری تحت برنامههای مختلف نرخ یادگیری (learning-rate schedules).
منبع: پست توسعهدهنده که جزئیات طراحی بهینهساز و نتایج تجربی را شرح میدهد.