‏چرا ران‌های پری‌ترینینگ شکست می‌خورن؟

آموزش مدل‌های زبانی بزرگ یه عملیات شکننده‌ست که چیزهای زیادی می‌تونه توش خراب بشه. دواکش پاتل در این یادداشت دو دسته اصلی از مشکلات رو معرفی می‌کنه: نقض کاوزالیتی (causality) و اضافه شدن بایاس به گرادیان‌ها.

در معماری‌های MoE (Mixture of Experts) که چند «متخصص» موازی دارن، روش expert choice routing یه مشکل جدی داره: وقتی هر اکسپرت خودش انتخاب می‌کنه که کدوم توکن‌ها رو پردازش کنه، ممکنه تخصیص یه توکن زودتر به یه اکسپرت به توکن‌های بعدی‌تر بستگی داشته باشه. این یعنی مدل در آموزش از اطلاعاتی استفاده می‌کنه که در استنتاج واقعی دسترسی بهشون نداره. شایعه‌ست که همین مشکل دلیل ضعف Llama 4 بوده. مشکل مشابهی به اسم token dropping هم وجود داره که گفته می‌شه Gemini 2 Pro باهاش دست و پنجه نرم کرده.

در بحث بایاس، یه مثال جالب از آموزش اولیه GPT-4 مطرح میشه: استفاده از FP16 در عملیات‌های collective مثل all-reduce باعث می‌شد که وقتی هزاران گرادیان کوچیک روی هم جمع می‌شن، به خاطر محدودیت دقت این فرمت اعداد، مقدار نهایی تا ۱۰ برابر خطا داشته باشه. پیدا کردن چنین باگی در یه سیستم توزیع‌شده عظیم، کار فوق‌العاده سختیه.

بخش دوم یادداشت به موازی‌سازی پری‌ترینینگ می‌پردازه. فرمول پایه اینه که کل FLOPs آموزش برابر 6 × N × D هست؛ N تعداد پارامترها و D تعداد توکن‌هاست. از اونجایی که یه GPU به تنهایی کافی نیست، باید بار رو تقسیم کرد. روش پیش‌فرض FSDP (Fully Sharded Data Parallel) اینه که هر GPU فقط ۱/N از پارامترهای هر لایه رو نگه می‌داره و قبل از پردازش هر لایه، وزن‌ها رو از بقیه جمع‌آوری می‌کنه. مزیت بزرگش اینه که می‌شه ارتباطات و محاسبات رو همزمان انجام داد چون وزن‌ها به خروجی لایه قبلی وابسته نیستن.

اما FSDP هم محدودیت داره. با افزایش تعداد GPU، زمان محاسبه کاهش پیدا می‌کنه ولی زمان ارتباطات ثابت می‌مونه، پس در یه نقطه‌ای گلوگاه ارتباطات می‌شه. علاوه بر این، چون هر GPU حداقل یه سکوانس کامل رو پردازش می‌کنه و attention درون یه سکوانس محاسبه میشه، اگه batch size محدود باشه، نمی‌شه از GPU های بیشتری استفاده کرد. اینجاست که pipeline parallelism وارد می‌شه؛ لایه‌های مختلف مدل روی GPU های مختلف قرار می‌گیرن. مشکل اصلیش اینه که در ابتدا و انتهای هر batch، بخشی از GPU ها بیکار می‌مونن (pipeline bubble) و این کارایی رو کاهش میده.

نکات کلیدی:

• نقض کاوزالیتی در روش‌های routing مثل expert choice یکی از مهم‌ترین دلایل شکست ران‌های پری‌ترینینگه
• بایاس خطرناک‌تر از واریانسه چون انباشته میشه، نه اینکه اوسط بگیره
• FSDP روش پیش‌فرض موازی‌سازیه و ارتباطاتش رو می‌شه با محاسبات هم‌پوشانی کرد
• pipeline parallelism علاوه بر bubble، محدودیت‌های معماری هم اعمال می‌کنه و سرعت تحقیق رو کاهش میده
• احتمالاً مشکلات جدید در هر مقیاس جدیدی ظهور می‌کنن و راه‌حل یکباره‌ای وجود نداره