‏موازی‌سازی مدل‌های بزرگ رو این‌بار با ویجت‌های زنده یاد بگیر

این صفحه یه بازنویسی تعاملیه از یکی از فصل‌های کتاب آنلاین «How to Scale Your Model» نوشتهٔ تیم Jax-ML؛ فصلی که دربارهٔ اینه وقتی می‌خوای یه مدل بزرگ رو رو صدها یا هزاران تراشه آموزش بدی، چطوری باید کارو بین تراشه‌ها تقسیم کنی. نویسندهٔ این نسخه به همون فرمول‌های کتاب اصلی، ویجت‌های زنده اضافه کرده که می‌شه باهاشون بازی کرد و اثر تغییر هر پارامتر رو دید.

ایدهٔ مرکزی مقاله اینه که رو هر تراشه، دو تا «ساعت» هم‌زمان کار می‌کنن: یکی ساعت محاسبه (چقدر طول می‌کشه واحد ضرب ماتریسی رو تراشه، عملیات ریاضی رو انجام بده) و یکی ساعت شبکه (چقدر طول می‌کشه داده‌ها بین تراشه‌ها رد و بدل بشه). اگه پیاده‌سازی درست باشه، این دو تا هم‌زمان اجرا می‌شن و زمان کل برابر بزرگ‌ترین‌شونه، نه مجموع‌شون؛ ولی اگه شبکه کندتر از محاسبه باشه، بخشی از تراشه بی‌کار می‌مونه.

مقاله پنج روش رایج موازی‌سازی رو مرور می‌کنه: data parallelism (که فقط داده رو بین تراشه‌ها تقسیم می‌کنه و وزن‌ها کامل رو هر تراشه می‌مونن)، FSDP یا ZeRO (که وزن‌ها و optimizer رو هم شارد می‌کنه تا حافظه کمتر مصرف بشه)، tensor parallelism (که خود لایه‌های مدل رو بین تراشه‌ها می‌شکنه)، expert parallelism (مخصوص مدل‌های Mixture-of-Experts) و در نهایت اشاره‌ای کوتاه به pipeline parallelism. برای هر کدوم نشون می‌ده دقیقاً کِی حجم ارتباطات از حجم محاسبات جلو می‌زنه و گلوگاه می‌شه.

برای مثال، تو data parallelism ساده، هر تراشه باید حداقل یه دسته (batch) به اندازهٔ نسبت سرعت محاسبه به پهنای باند شبکه پردازش کنه تا گلوگاه شبکه نشه؛ رو یه TPUv5p این عدد حدود ۲۵۵۰ توکن به ازای هر تراشه‌ست. تو مدل‌های Mixture-of-Experts این آستانه با ضریب E/k (نسبت کل اکسپرت‌ها به اکسپرت‌های فعال) بزرگ‌تر هم می‌شه، چون گرادیان همهٔ اکسپرت‌ها رد و بدل می‌شه ولی محاسبات فقط رو اکسپرت‌های فعال انجام می‌گیره.

نکات کلیدی:

• نسخهٔ تعاملیِ فصل موازی‌سازی از کتاب Jax-ML Scaling Book، با ویجت‌های زنده رو اعداد واقعی سخت‌افزار
• ایدهٔ اصلی: مقایسهٔ هم‌زمان «ساعت محاسبه» و «ساعت شبکه» رو هر تراشه (مدل roofline)
• پنج روش بررسی می‌شه: data parallelism، FSDP/ZeRO، tensor parallelism، expert parallelism و اشاره‌ای به pipeline parallelism
• رو TPUv5p، آستانهٔ گلوگاه شبکه برای data parallelism ساده حدود ۲۵۵۰ توکن به ازای هر تراشه‌ست؛ تو مدل‌های MoE این عدد با ضریب E/k بیشتر می‌شه