صاف‌کردن قله‌های حافظه در آموزش مدل‌های MoE
هوش مصنوعیخلاصهٔ کامل‌تر
به گفتهٔ نویسنده‌های این مقالهٔ arXiv، وقتی یه مدل Mixture-of-Experts (یعنی مدلی که به جای یک شبکهٔ بزرگ چند تا «متخصص» داره و هر توکن فقط به بعضی‌هاشون مسیردهی میشه) رو با کانتکست طولانی یا بچ بزرگ آموزش می‌دن، کافیه پیک حافظهٔ فقط یکی از اجزا از حافظهٔ کارت رد بشه تا کل آموزش بخوابه. پس چیزی که باید مهار بشه میانگین مصرف حافظه نیست، خود پیک‌هاست، اونم همه‌شون با هم.
چهار پیک هست که تو پلن‌های رایج موازی‌سازی هیچ سقفی ندارن، و هرکدوم جور دیگه‌ای بزرگ می‌شن: expert dispatch با ماتریس مسیردهی، vocabulary projection با تعداد توکن ضربدر اندازهٔ واژگان، مرزهای gradient checkpoint با عمق ضربدر طول دنباله، و optimizer state با تعداد پارامتر. کدومشون زودتر پر میشه به مدل و طول کانتکست و تعداد کارت بستگی داره، برای همین پایین‌آوردن بزرگ‌ترین پیک فقط نوبت رو می‌ده به پیک بعدی.
راه‌حل چهار تا زمان‌بندیه که working set روی GPU رو از همون لحظهٔ اجرا ثابت نگه می‌دارن. PipelinedLLEP سهم توکن هر مبدأ توی یه chunk از dispatch رو سقف می‌ذاره، Ring-DTP فعال‌سازی‌ها یا شاردهای وزن رو دور یه حلقه می‌چرخونه و هر بلوک از logitها رو با log-sum-exp آنلاین جمع می‌کنه، SCO تنها تنسور طولانی‌عمر هر مرز checkpoint رو تو حافظهٔ CPU نگه می‌داره، و OffloadStreamAdamW آپدیت سریالی Adam روی CPU رو به یه خط لولهٔ سطلی تبدیل می‌کنه.
نویسنده‌ها تأکید می‌کنن هر چهار روش فقط ترتیب و دانه‌بندی محاسبه و جابه‌جایی داده رو عوض می‌کنن، پس loss و گرادیان دقیقاً همونی می‌مونه که بود. تو تست‌های جداگانه، پیک dispatch تا ۵۹.۳٪ و پیک vocabulary projection تا ۸۶.۶٪ پایین اومده و قدم آپدیت optimizer آفلودشده ۲.۰۵ برابر سریع‌تر شده. با هم روی مدل‌های ۱۲۰ تا ۶۶۷ میلیارد پارامتری، آموزش تا کانتکست ۱ میلیون توکن جلو رفته، یعنی ۸ تا ۳۲ برابر بردِ یه baseline تیون‌شدهٔ FSDP2 و تا ۱۰.۴ برابر throughput اون.
نکات کلیدی:
- چهار پیک بی‌سقف: expert dispatch، vocabulary projection، مرزهای gradient checkpoint و optimizer state
- PipelinedLLEP پیک dispatch رو تا ۵۹.۳٪ و Ring-DTP پیک vocabulary projection رو تا ۸۶.۶٪ کم کرده
- SCO تنسور طولانی‌عمر هر مرز checkpoint رو به حافظهٔ CPU می‌بره؛ OffloadStreamAdamW قدم Adam رو ۲.۰۵ برابر سریع کرده
- روی مدل‌های ۱۲۰B تا ۶۶۷B: کانتکست ۱M توکن، ۸ تا ۳۲ برابر برد و تا ۱۰.۴ برابر throughput نسبت به FSDP2
- loss و گرادیان دست‌نخورده می‌مونه، چون فقط ترتیب محاسبه عوض شده
