آموزش یه مدل MoE روی دو کلاستر با ۹۰۰۰ مایل فاصله
هوش مصنوعیخلاصهٔ کامل‌تر
نویسنده می‌گه مشکل اصلی آموزش مدل‌های frontier کمبود چیپ نیست، کمبود برق، تجهیزات شبکهٔ برق و زمینِ مناسبه. ترانسفورمرهای بزرگ بیشتر از سه سال طول می‌کشه برسن و Goldman Sachs پیش‌بینی می‌کنه فقط ۵۰ تا ۶۰ درصد ظرفیت دیتاسنترهای برنامه‌ریزی‌شدهٔ آمریکا سر وقت راه بیفته. نتیجه‌ش اینه که GPUها تو compute islandها پخش شدن، یعنی کلاسترهای چندصد تا چندهزار GPUیی که هر کدوم جدا هستن. ولی بیشتر stackهای آموزش فرض می‌کنن همه‌چی تو یه کلاستر بزرگ و یکدست هست.
به گفتهٔ مقاله، آموزش با ارتباط کم ایدهٔ جدیدی نیست و ریشه‌ش به دههٔ ۲۰۱۰ برمی‌گرده: Local SGD (یعنی چند قدم آموزش محلی قبل از هر تبادل)، فشرده‌سازی گرادیان، gossip (هر نود فقط با چند همسایه حرف می‌زنه) و روش‌های مقاوم در برابر خرابی. از ۲۰۲۳ به بعد روش‌هایی مثل DiLoCo نشون دادن که با ۸ worker میشه به کیفیت آموزش کاملاً همگام رسید، با ۵۰۰ برابر ارتباط کمتر. نویسنده دو خانواده رو جدا می‌کنه: data-parallel که هر سایت یه کپی کامل از مدل داره، و model-parallel که مدل بین شرکت‌کننده‌ها تقسیم میشه.
تو آزمایش اصلی، تیم از ۱۶۰ تا B300 تو آمریکا و ۱۶۰ تا GB300 تو آسیا استفاده کرد. مدل یه MoE (یعنی مدلی که برای هر توکن فقط چند تا از «متخصص»هاش فعال می‌شن) با حدود ۳۲ میلیارد پارامتر، ۲۵۶ expert و مسیریابی top-8 بود. دو کلاستر batch size و روش موازی‌سازی متفاوتی داشتن. مدل تکه‌تکه و با فاصله تو هر پنجرهٔ ۵۰ قدمی همگام می‌شد، هر سایت یه نسخه از Distributed Shampoo رو به‌عنوان optimizer محلی داشت، مقادیر با stochastic rounding و error feedback فشرده می‌شدن و آمار routing مربوط به MoE تو هر قدم جدا همگام می‌شد.
لایهٔ شبکه هم دوباره نوشته شد تا روی tail latency تمرکز کنه، یعنی کندترین انتقال‌ها که کل آپدیت رو معطل می‌کنن. ترافیک تو زمان پخش شد، اتصال‌ها گرم نگه داشته شدن، payloadها به تکه‌های مستقل تقسیم شدن که هر اتصال سالمی بتونه برشون داره، و فقط تکه‌های گم‌شده دوباره درخواست شدن. با همون مقدار compute، اجرای دوسایتی به حدود ۰.۲۵۲۵ bits per byte رسید در برابر ۰.۲۵۴۹ برای آموزش استاندارد، یعنی حدود ۱٪ بهتر.
تو آزمایش دوم، ۶۴ تا GB300 به هشت سلول هشت‌تایی تقسیم شدن که هر کدوم کپی خودش از مدل رو آموزش می‌داد. به‌جای all-reduce (یعنی میانگین‌گیری هم‌زمان بین همه)، سلول‌ها تکه‌های مدل رو دوبه‌دو عوض می‌کردن و بعد از سه دور، آپدیت هر سلول به همهٔ هشت‌تا می‌رسید. این روش butterfly mixing از Zhao و Canny هست. هر تکه شش بار تو هر ۵۰ قدم تبادل می‌شد و سلول دو قدم بعد نتیجهٔ جفتش رو با وزن مساوی با وزن‌های خودش ترکیب می‌کرد. کیفیت آموزش با baseline برابر شد.
نویسنده معتقده اگه آموزش بتونه روی چند کلاستر کوچیک‌تر پخش بشه، خیلی از ظرفیت محاسباتی دنیا برای آموزش frontier قابل استفاده میشه، حتی ظرفیت ارزون و preemptible و نسل‌های قدیمی‌تر GPU. اون حتی به آموزش تو مدار زمین اشاره می‌کنه، مثل Project Suncatcher گوگل. ولی هشدار هم می‌ده که وقتی آموزش بین چند ارائه‌دهنده پخش بشه، یه سایت خراب یا نفوذشده می‌تونه بی‌صدا مدل رو خراب کنه و ردیابیش سخت‌تره.
نکات کلیدی:
- مدل MoE با حدود ۳۲ میلیارد پارامتر، ۲۵۶ expert و مسیریابی top-8 روی دو کلاستر با ۹۰۰۰ مایل فاصله آموزش داده شد.
- کلاسترها GPUهای متفاوت داشتن: ۱۶۰ تا B300 تو آمریکا و ۱۶۰ تا GB300 تو آسیا، با لینک دوطرفهٔ ۲۵ گیگابیتی.
- نتیجهٔ دوسایتی ۰.۲۵۲۵ bits per byte بود در برابر ۰.۲۵۴۹ برای آموزش استاندارد.
- همگام‌سازی تکه‌به‌تکه و با فاصله تو پنجره‌های ۵۰ قدمی انجام شد، با الهام از Streaming DiLoCo و Decoupled DiLoCo.
- butterfly mixing با هشت learner و تبادل دوبه‌دو، بدون all-reduce به کیفیت baseline رسید.
- ریسک امنیتی: یه سایت نفوذشده می‌تونه بی‌صدا مدل رو دستکاری کنه.
