آموزش توزیع‌شده مدل‌های زبانی: FSDP، DeepSpeed و موازی‌سازی سه‌بعدی

وقتی مدل‌های زبانی از چند صد میلیون پارامتر عبور می‌کنند، نگهداری وزن‌ها، گرادیان‌ها و وضعیت‌های بهینه‌ساز فراتر از گنجایش حافظه‌ی یک GPU معمولی می‌شود. در چنین شرایطی، آموزش مدل دیگر روی یک کارت گرافیک جا نمی‌شود و استفاده از تکنیک‌های آموزش توزیع‌شده (Distributed Training) تنها راهکار عملی است.

تهیه‌ی GPU قدرتمند یا اجاره‌ی آن از سرویس‌های ابری خارجی برای کاربران ایرانی دشوار و پرهزینه است و کارت‌های بانکی ایرانی نیز امکان پرداخت مستقیم به پلتفرم‌های بین‌المللی را ندارند. در نتیجه، پیش از صرف هزینه یا دسترسی به زیرساخت واقعی، باید بدانید حافظه چطور مصرف می‌شود و مدل را با چه استراتژی‌هایی بهینه‌سازی کنید.

در این درس یاد می‌گیرید که چگونه بودجه‌ی دقیق حافظه را پیش از شروع آموزش محاسبه کنید، تفاوت استراتژی‌های موازی‌سازی داده، تنسور و پایپ‌لاین را بشناسید و با مکانیزم‌هایی مثل FSDP و DeepSpeed ZeRO مدل‌های چند میلیاردی را حتی با منابع محدودتر سخت‌افزاری مدیریت کنید.

مدل 124M شما روی یک کارت گرافیک آموزش دید. حالا 7 میلیارد پارامتر را امتحان کنید؛ مدل در حافظه جا نمی‌شود. پردازش داده‌ها روی یک ماشین هفته‌ها زمان می‌برد. در مقیاس‌های بزرگ، آموزش توزیع‌شده (distributed training) یک گزینه‌ی اختیاری نیست؛ تنها راه پیش رو است.

اهداف یادگیری

• توضیح سه نوع موازی‌سازی (داده، تنسور، پایپ‌لاین) و تشخیص زمان نیاز به هر کدام بر اساس اندازه‌ی مدل و اندازه‌ی کلاستر
• پیاده‌سازی آموزش با موازی‌سازی داده با PyTorch DDP به همراه همگام‌سازی گرادیان روی چند کارت گرافیک
• محاسبه‌ی بودجه‌ی حافظه برای اندازه‌ی مشخصی از مدل (وزن‌ها + وضعیت‌های بهینه‌ساز + گرادیان‌ها + اکتیویشن‌ها) جهت تعیین حداقل سخت‌افزار مورد نیاز
• پیکربندی مراحل FSDP یا DeepSpeed ZeRO برای شارد کردن (shard) وضعیت‌های مدل روی کارت‌های گرافیک و جا دادن مدل‌هایی که از حافظه‌ی یک کارت گرافیک فراتر می‌روند

صورت مسئله

یک مدل با 7B پارامتر در فرمت FP16 تنها برای وزن‌ها به 14GB حافظه نیاز دارد. بهینه‌ساز Adam دو نسخه‌ی دیگر از هر پارامتر را ذخیره می‌کند (تخمین گشتاورهای اول و دوم). این یعنی 28GB دیگر. گرادیان‌ها هنگام پس‌انتشار خطا 14GB دیگر اضافه می‌کنند. قبل از اینکه حتی یک اکتیویشن (activation) ذخیره شود، به 56GB رسیده‌اید.

یک NVIDIA A100 معادل 80GB حافظه دارد.

56GB از 80GB مصرف شده است. این یعنی 24GB برای اکتیویشن‌ها باقی می‌ماند؛ مقادیر میانی محاسبه‌شده در گذر رو به جلو که باید برای پس‌انتشار خطا در حافظه زنده بمانند. برای یک دنباله‌ی 2048 توکنی در مدلی با بعد 4096، اکتیویشن‌های یک لایه حدود 64MB حافظه مصرف می‌کنند. با 32 لایه، به ازای هر نمونه به 2GB نیاز دارید. اندازه‌ی دسته (batch size) برابر با 8 به 16GB فضا نیاز دارد. شما 24GB در اختیار دارید؛ اندازه‌ی دسته‌ی 12 سیستم را با خطای کمبود حافظه متوقف می‌کند.

حالا 70B پارامتر را امتحان کنید. فقط وزن‌ها: 140GB در FP16. روی یک کارت گرافیک جا نمی‌شود. حداقل به 2 کارت A100 (2 × 80GB = 160GB) فقط برای نگهداری وزن‌ها نیاز دارید. وضعیت‌های بهینه‌ساز و گرادیان‌ها را هم اضافه کنید تا ببینید به حافظه‌ی بسیار بیشتری نیاز خواهید داشت: حداقل +3 کارت گرافیک، و در واقعیت بسته به استراتژی شاردینگ، بین 8 تا 16 کارت.

مدل Llama 3 405B روی 16,384 کارت گرافیک NVIDIA H100 آموزش داده شد. هزینه‌ی رایانشی این اجرای آموزشی حدود 100 میلیون دلار تخمین زده می‌شود. DeepSeek V3 با هوشمندی در معماری (تلفیق متخصصان یا Mixture of Experts به این معنی است که به ازای هر توکن تنها کسری از پارامترها فعال می‌شوند) و بالا بردن بهره‌وری آموزش، مدلی هم‌رده را با هزینه‌ای در حدود 5.6 میلیون دلار آموزش داد.

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

مفاهیم پایه

چرا توزیع‌شدگی الزامی است

محاسبات حافظه برای مدل‌های واقعی به شرح زیر است. هر عدد بر پایه‌ی محاسبه به دست آمده، نه تخمین.

ستون «وضعیت‌های Adam» کمرشکن است. Adam برای هر پارامتر یک میانگین متحرک (m) و یک واریانس متحرک (v) ذخیره می‌کند که هر دو در فرمت FP32 هستند. برای یک مدل 70B، این مقدار برابر است با: 70B × 4 بایت × 2 = 560GB. بهینه‌ساز به تنهایی به هفت کارت A100 نیاز دارد.

یک H100 منفرد دارای 80GB حافظه است. Llama 3 405B حداقل به 61 کارت H100 فقط برای نگهداری وزن‌ها، بهینه‌ساز و گرادیان‌ها نیاز دارد. اکتیویشن‌ها را اضافه کنید تا این عدد باز هم بزرگ‌تر شود. متا از 16,384 کارت گرافیک استفاده نکرد چون دلش می‌خواست؛ از سر اجبار این کار را کرد.

موازی‌سازی داده

ساده‌ترین استراتژی توزیع‌شده است. کل مدل را روی N کارت گرافیک کپی کنید. هر دسته‌ی آموزشی را به N بخش مساوی تقسیم کنید. هر کارت گرافیک یک گذر رو به جلو و یک گام رو به عقب روی شارد داده‌ی خود اجرا می‌کند. پس از گام رو به عقب، گرادیان‌ها را بین تمام کارت‌های گرافیک میانگین‌گیری کنید. هر کارت گرافیک نسخه‌ی وزن‌های خود را با همان گرادیان‌های میانگین‌گیری‌شده به‌روزرسانی می‌کند و همه‌ی کپی‌ها همگام می‌مانند.

مزیت: مقیاس‌پذیری خطی نرخ پردازش. با N کارت گرافیک، در هر گام N برابر داده‌ی بیشتری پردازش می‌شود. تبادل داده فقط به میانگین‌گیری گرادیان محدود می‌شود که با محاسبات هم‌پوشانی دارد.

عیب: هر کارت گرافیک یک نسخه‌ی کامل از مدل، وضعیت‌های بهینه‌ساز و گرادیان‌ها را نگه می‌دارد. برای یک مدل 70B، هر کارت به 840GB حافظه نیاز دارد. موازی‌سازی داده حافظه‌ی مصرفی به ازای هر کارت گرافیک را کم نمی‌کند؛ فقط زمان آموزش را کاهش می‌دهد.

محاسبات: اندازه‌ی موثر دسته = اندازه‌ی دسته به ازای هر کارت گرافیک × N. برای N=64 کارت گرافیک با دسته‌ی 16تایی روی هر کارت، دسته‌ی موثر 1,024 خواهد بود. Llama 3 از اندازه‌ی دسته‌ی موثر 16 میلیون توکن در هر گام استفاده کرد.

موازی‌سازی تنسور

لایه‌های مجزا را بین کارت‌های گرافیک تقسیم کنید. یک ضرب ماتریسی منفرد بین کارت‌های گرافیک خرد می‌شود و هر کدام بخشی از نتیجه را محاسبه می‌کنند.

یک ماتریس وزن با ابعاد (8192, 8192) را در یک لایه‌ی پیش‌خور (feedforward) در نظر بگیرید. با موازی‌سازی تنسور 4طرفه، هر کارت گرافیک یک شارد با ابعاد (8192, 2048) را نگه می‌دارد. هر کارت ورودی را در شارد خود ضرب می‌کند و یک نتیجه‌ی جزئی به دست می‌آورد. این نتایج جزئی با هم ترکیب می‌شوند (از طریق all-reduce یا all-gather) تا خروجی کامل حاصل شود.

مزیت: کاهش حافظه‌ی لازم به ازای هر کارت گرافیک برای وزن‌های مدل. تقسیم یک مدل 70B بین 8 کارت گرافیک یعنی هر کارت به اندازه‌ی حدود 8.75B پارامتر وزن نگه می‌دارد.

عیب: نیاز به ارتباط بسیار سریع بین کارت‌های گرافیک پس از هر لایه دارد. عملیات all-reduce پس از هر ضرب ماتریسی تاخیر ایجاد می‌کند. این رویکرد روی NVLink (با پهنای باند 900 GB/s بین کارت‌های یک نود) عالی جواب می‌دهد، اما بین نودهای مختلف که با InfiniBand به هم متصل شده‌اند (با سرعت 400 Gb/s، حدود 50 GB/s) عملکرد ضعیفی دارد. موازی‌سازی تنسور تقریباً همیشه به درون یک نود منفرد (8 کارت گرافیک) محدود می‌شود.

کاربرد واقعی: Megatron-LM پیشگام موازی‌سازی تنسور بود. Llama 3 405B از موازی‌سازی تنسور 8طرفه درون هر نود استفاده می‌کند.

موازی‌سازی پایپ‌لاین

مدل را بر اساس لایه‌ها تقسیم کنید. کارت گرافیک 1 لایه‌های 1 تا 8 را اجرا می‌کند. کارت گرافیک 2 لایه‌های 9 تا 16 را اجرا می‌کند. کارت گرافیک 3 لایه‌های 17 تا 24 را اجرا می‌کند. کارت گرافیک 4 لایه‌های 25 تا 32 را اجرا می‌کند. داده‌ها در طول پایپ‌لاین جریان می‌یابند: کارت گرافیک 1 لایه‌های خود را محاسبه کرده و اکتیویشن‌ها را به کارت گرافیک 2 می‌فرستد؛ کارت 2 لایه‌های خود را محاسبه کرده و به کارت 3 می‌فرستد و به همین ترتیب.

مزیت: حداقل ارتباط بین کارت‌های گرافیک؛ صرفاً اکتیویشن‌های مرز لایه‌ها منتقل می‌شوند که در مقایسه با گرادیان‌ها یا وزن‌ها بسیار کوچک هستند. به دلیل نیاز اندک به پهنای باند، روی چند نود مجزا نیز به خوبی کار می‌کند.

عیب: حباب‌های پایپ‌لاین (pipeline bubbles). وقتی کارت گرافیک 4 در حال محاسبه‌ی گذر رو به جلو روی ریزدسته‌ی (micro-batch) شماره 1 است، کارت‌های 1 و 2 و 3 بیکار هستند (بخش خود را قبلاً جلو برده‌اند). در گام رو به عقب، این الگو معکوس می‌شود. در پایپ‌لاین ساده، میزان بهره‌وری کارت‌های گرافیک برای N مرحله‌ی پایپ‌لاین فقط 1/N است.

روش‌های GPipe و PipeDream مشکل حباب را با خرد کردن هر دسته به چند ریزدسته حل می‌کنند. کارت گرافیک 1 به محض اینکه گذر رو به جلو ریزدسته‌ی 1 را تمام کرد، کار روی ریزدسته‌ی 2 را آغاز می‌کند. این کار محاسبات را در تمام مراحل پایپ‌لاین هم‌پوشانی می‌دهد. با M ریزدسته و N مرحله، کسر حباب به (N-1)/M کاهش می‌یابد. اگر از M=16 ریزدسته با N=4 مرحله استفاده کنید، حباب به 3/16 = 18.75% زمان بیکاری کاهش پیدا می‌کند.

اف‌اس‌دی‌پی (FSDP): موازی‌سازی داده با شاردینگ کامل

روش FSDP مقیاس‌پذیری موازی‌سازی داده را با بهره‌وری حافظه در شاردینگ ترکیب می‌کند. به جای اینکه هر کارت گرافیک یک نسخه‌ی کامل از مدل را نگه دارد، هر کارت تنها 1/N از پارامترها، گرادیان‌ها و وضعیت‌های بهینه‌ساز را در خود جای می‌دهد.

پیش از گذر رو به جلو در هر لایه، FSDP یک عملیات all-gather اجرا می‌کند تا تمام پارامترها را از همه‌ی کارت‌های گرافیک در حافظه‌ی هر کارت جمع‌آوری کند. پس از گذر رو به جلو، هر کارت گرافیک پارامترهای غیرمحلی را دور می‌ریزد. در گام رو به عقب، all-gather دوباره اجرا می‌شود تا پارامترها برای محاسبه‌ی گرادیان بازسازی شوند. پس از گام رو به عقب، یک عملیات reduce-scatter شاردهای گرادیان را توزیع می‌کند تا هر کارت فقط 1/N از گرادیان‌ها را ذخیره کند.

محاسبات برای یک مدل 70B روی 8 کارت گرافیک:

بدون FSDP، جا دادن یک مدل 70B روی یک کارت گرافیک 80GB ناممکن است. با FSDP روی 8 کارت گرافیک، هر کارت 105GB مصرف می‌کند؛ دست نگه دارید، این مقدار هم هنوز جا نمی‌شود. برای رسیدن به زیر 80GB به ازای هر کارت، حداقل به 16 کارت گرافیک نیاز دارید یا اینکه باید FSDP را با چک‌پوینت‌گذاری اکتیویشن (محاسبه‌ی دوباره‌ی اکتیویشن‌ها در گام رو به عقب به جای ذخیره‌ی آن‌ها) ترکیب کنید.

هزینه‌ی ارتباطات به دلیل all-gather قبل از هر لایه، از موازی‌سازی داده‌ی معمولی بیشتر است. اما صرفه‌جویی در حافظه، اجراهایی را که قبلاً ناممکن بودند عملی می‌کند.

ابزار DeepSpeed ZeRO

ماژول ZeRO (بهینه‌ساز افزونگی صفر یا Zero Redundancy Optimizer) در DeepSpeed از نظر مفهومی با FSDP یکسان است اما به‌طور مستقل توسط مایکروسافت توسعه یافت. این ابزار سه مرحله (stage) تعریف می‌کند که در هر مرحله شاردینگ تهاجمی‌تر می‌شود:

مرحله‌ی ZeRO-3 معادل FSDP است. نام‌گذاری‌ها فرق دارد، ولی سازوکار یکی است. PyTorch بعد از اینکه DeepSpeed کارآمدی این ایده را اثبات کرد، FSDP را به صورت بومی اضافه نمود.

ابزار DeepSpeed همچنین ZeRO-Offload (انتقال وضعیت‌های بهینه‌ساز به رم CPU که ارزان‌تر و حجیم‌تر است) و ZeRO-Infinity (انتقال به حافظه‌های NVMe SSD) را معرفی کرد. این قابلیت‌ها سرعت پردازش را فدای ظرفیت حافظه می‌کنند؛ عملیات منتقل‌شده کندتر اجرا می‌شوند اما حافظه‌ی کارت گرافیک را آزاد می‌کنند.

آموزش با دقت ترکیبی

آموزش مدرن از چند فرمت ممیز شناور به صورت همزمان استفاده می‌کند:

• گذر رو به جلو: فرمت FP16 یا BF16 (16 بیتی). نصف حافظه‌ی FP32. ضرب‌های ماتریسی روی هسته‌های تنسور (Tensor Core) دو برابر سریع‌تر اجرا می‌شوند.
• وزن‌های اصلی (Master weights): فرمت FP32 (32 بیتی). توسط بهینه‌ساز برای حفظ دقت عددی هنگام به‌روزرسانی وزن‌ها نگهداری می‌شوند.
• مقیاس‌بندی تابع زیان: ضرب کردن تابع زیان (loss) در یک عدد ثابت بزرگ پیش از گام رو به عقب برای جلوگیری از کم‌ریز عددی گرادیان‌های FP16 به صفر. قبل از گام بهینه‌ساز، مقادیر بر همان عدد ثابت تقسیم می‌شوند.

فرمت BF16 (مخفف Brain Float 16) همان محدوده‌ی توان FP32 را دارد (8 بیت توان) ولی دقت کمتری ارائه می‌دهد (7 بیت مانتیس در برابر 23 بیت در FP32). این فرمت به‌ندرت به مقیاس‌بندی تابع زیان نیاز پیدا می‌کند چون می‌تواند همان دامنه‌ی مقادیر را پوشش دهد. در مقابل، FP16 دارای 5 بیت توان و 10 بیت مانتیس است؛ می‌تواند مقادیر ریزتر را نمایش دهد اما در مقادیر بسیار بزرگ یا کوچک دچار سرریز یا کم‌ریز عددی می‌شود.

پردازنده‌های TPU گوگل ذاتاً از BF16 استفاده می‌کنند. کارت‌های A100 و H100 شرکت NVIDIA از هر دو فرمت FP16 و BF16 پشتیبانی می‌کنند. صنعت هوش مصنوعی تا حد زیادی به سمت BF16 کوچ کرده است، زیرا دردسرهای مقیاس‌بندی تابع زیان را برطرف می‌کند.

مقایسه‌ی حافظه برای یک مدل 7B:

دقت ترکیبی در این مدل 28GB صرفه‌جویی می‌کند. وضعیت‌های بهینه‌ساز در هر صورت روی FP32 باقی می‌مانند؛ بخش عمده‌ی حافظه همین‌جا می‌رود.

کتابخانه‌ی Megatron-LM و موازی‌سازی سه‌بعدی

آموزش‌های بزرگ در دنیای واقعی هر سه نوع موازی‌سازی را با هم ترکیب می‌کنند:

• موازی‌سازی داده در بین گروه‌هایی از نودها (برای افزایش اندازه‌ی دسته)
• موازی‌سازی تنسور درون یک نود (تقسیم لایه‌ها بین 8 کارت گرافیک)
• موازی‌سازی پایپ‌لاین بین نودها (تقسیم دسته‌های لایه در میان ماشین‌های مختلف)

آموزش Llama 3 405B روی 16,384 کارت H100:

• موازی‌سازی تنسور 8طرفه در هر نود (8 کارت گرافیک در هر نود)
• موازی‌سازی پایپ‌لاین 16طرفه بین نودها (16 مرحله‌ی پایپ‌لاین)
• موازی‌سازی داده‌ی 128طرفه در بعد باقی‌مانده (16,384 / 8 / 16 = 128)

این تفکیک سه‌بعدی (8 × 16 × 128 = 16,384) روشی است که برای مقیاس‌بندی تا هزاران کارت گرافیک استفاده می‌شود. هر کارت گرافیک یک شارد داده‌ی متفاوت را می‌بیند (موازی داده)، برشی از هر لایه را نگه می‌دارد (موازی تنسور)، و مجموعه‌ی متفاوتی از لایه‌ها را پردازش می‌کند (موازی پایپ‌لاین).

مدل DeepSeek V3 رویکرد دیگری پیش گرفت. معماری تلفیق متخصصان آن‌ها تنها 37B از 671B پارامتر را به ازای هر توکن فعال می‌کند. یعنی هر کارت گرافیک فقط باید پارامترهای فعال را محاسبه کند (و اکتیویشن‌های آن‌ها را نگه دارد). آن‌ها مدل را روی 2,048 کارت گرافیک H800 آموزش دادند — کمتر از یک‌هشتم تعداد کارت‌های گرافیک متا — با هزینه‌ی 5.6 میلیون دلار در برابر رقم تخمینی 100 میلیون دلاری متا.

پیاده‌سازی از صفر

گام 1: شبیه‌سازی موازی‌سازی داده

یک دسته را بین کارت‌های گرافیک شبیه‌سازی‌شده تقسیم کنید. هر کارت یک گذر رو به جلو روی شارد خود اجرا می‌کند. سپس «گرادیان‌ها» را میانگین‌گیری کنید (ما آن‌ها را به صورت مقادیر تابع زیان شبیه‌سازی می‌کنیم).

import numpy as np

def simulate_data_parallelism(data, num_gpus, model_fn):
    batch_size = len(data)
    shard_size = batch_size // num_gpus
    remainder = batch_size % num_gpus

    gpu_losses = []
    gpu_gradients = []

    offset = 0
    for gpu_id in range(num_gpus):
        extra = 1 if gpu_id < remainder else 0
        shard = data[offset:offset + shard_size + extra]
        offset += shard_size + extra

        loss, grad = model_fn(shard)
        gpu_losses.append(loss)
        gpu_gradients.append(grad)

    avg_loss = np.mean(gpu_losses)
    avg_gradient = np.mean(gpu_gradients, axis=0)

    return avg_loss, avg_gradient

عملیات all-reduce (میانگین‌گیری گرادیان‌ها) تنها تبادل داده در موازی‌سازی داده است. در عمل، این کار با کتابخانه‌ی NCCL روی کارت‌های گرافیک NVIDIA انجام می‌شود که الگوریتم ring all-reduce را پیاده کرده است: هر کارت گرافیک 1/N از گرادیان‌هایش را به همسایه‌اش می‌فرستد، 1/N از همسایه‌ی دیگر دریافت می‌کند و بعد از N-1 گام، همه‌ی کارت‌ها میانگین کامل را دارند. حجم کل تبادل داده: 2 × gradient_size × (N-1)/N که برای Nهای بزرگ به سمت 2 برابر اندازه‌ی گرادیان میل می‌کند.

گام 2: شبیه‌سازی موازی‌سازی تنسور

یک ماتریس وزن را بین کارت‌های گرافیک تقسیم کنید. هر کارت گرافیک یک ضرب ماتریسی جزئی را محاسبه می‌کند. سپس نتایج را ترکیب کنید.

def simulate_tensor_parallelism(input_data, weight_matrix, num_gpus):
    d_in, d_out = weight_matrix.shape
    assert d_out % num_gpus == 0, f"d_out {d_out} not divisible by num_gpus {num_gpus}"
    shard_size = d_out // num_gpus

    partial_results = []
    for gpu_id in range(num_gpus):
        start = gpu_id * shard_size
        end = start + shard_size
        weight_shard = weight_matrix[:, start:end]

        partial = input_data @ weight_shard
        partial_results.append(partial)

    full_output = np.concatenate(partial_results, axis=-1)

    direct_output = input_data @ weight_matrix
    error = np.abs(full_output - direct_output).max()

    return full_output, error

خطا باید دقیقاً صفر (یا در حد اپسیلون ماشین) باشد. موازی‌سازی تنسور از نظر ریاضی کاملاً دقیق است؛ خروجی آن دقیقاً همان نتیجه‌ی ضرب کامل ماتریس روی یک کارت گرافیک است. این تقسیم در امتداد بعد خروجی انجام شده، بنابراین هر کارت گرافیک دسته‌ی متفاوتی از ستون‌ها را تولید می‌کند و الحاق (concatenation) آن‌ها نتیجه‌ی کامل را بازسازی می‌کند.

برای لایه‌های خطی موازی-ستونی (تقسیم بعد خروجی)، خروجی‌ها را الحاق می‌کنید. برای لایه‌های موازی-سطری (تقسیم بعد ورودی)، خروجی‌ها را جمع می‌زنید. در بخش FFN یک ترنسفورمر، لایه‌ی خطی اول (توسعه) از موازی-ستونی و لایه‌ی خطی دوم (انقباض) از موازی-سطری استفاده می‌کند. این کار مانع از اجرای یک all-reduce اضافی بین دو لایه می‌شود.

گام 3: شبیه‌سازی موازی‌سازی پایپ‌لاین

لایه‌های یک مدل را بین کارت‌های گرافیک مجازی تقسیم کنید. مشکل حباب را نمایش دهید که در آن مراحل اولیه در زمان پردازش مراحل بعدی، بیکار می‌مانند.

def simulate_pipeline_parallelism(num_layers, num_stages, num_microbatches):
    layers_per_stage = num_layers // num_stages

    timeline = {}
    clock = 0

    for mb in range(num_microbatches):
        for stage in range(num_stages):
            start_time = max(
                timeline.get((stage, mb - 1, "fwd"), (0, 0))[1] if mb > 0 else 0,
                timeline.get((stage - 1, mb, "fwd"), (0, 0))[1] if stage > 0 else 0,
            )
            end_time = start_time + layers_per_stage
            timeline[(stage, mb, "fwd")] = (start_time, end_time)

    last_fwd_end = max(v[1] for v in timeline.values())

    for mb in range(num_microbatches - 1, -1, -1):
        for stage in range(num_stages - 1, -1, -1):
            deps = [last_fwd_end]
            if mb < num_microbatches - 1 and (stage, mb + 1, "bwd") in timeline:
                deps.append(timeline[(stage, mb + 1, "bwd")][1])
            if stage < num_stages - 1 and (stage + 1, mb, "bwd") in timeline:
                deps.append(timeline[(stage + 1, mb, "bwd")][1])
            start_time = max(deps)
            end_time = start_time + layers_per_stage
            timeline[(stage, mb, "bwd")] = (start_time, end_time)

    total_time = max(v[1] for v in timeline.values())
    compute_time = num_microbatches * num_stages * layers_per_stage * 2
    bubble_fraction = 1.0 - compute_time / (total_time * num_stages)

    return timeline, total_time, bubble_fraction

با 4 مرحله و 1 ریزدسته، کسر حباب 75% است؛ در هر لحظه سه کارت از چهار کارت گرافیک بیکار هستند. با 16 ریزدسته، این رقم به حدود 19% می‌رسد. بهای حذف حباب‌ها مصرف حافظه است: باید اکتیویشن‌های تمام ریزدسته‌های در حال پردازش را به صورت همزمان ذخیره کنید.

گام 4: ماشین‌حساب محاسبه‌ی حافظه

نیازمندی‌های دقیق حافظه را برای آموزش هر اندازه‌ای از مدل محاسبه کنید.

def memory_calculator(
    params_billions,
    precision_bytes=2,
    optimizer="adam",
    num_gpus=1,
    sharding="none",
    sequence_length=2048,
    batch_size_per_gpu=1,
    hidden_dim=None,
    num_layers=None,
):
    params = params_billions * 1e9

    weight_memory = params * precision_bytes

    if optimizer == "adam":
        optimizer_memory = params * 4 * 2
    elif optimizer == "sgd":
        optimizer_memory = params * 4
    else:
        optimizer_memory = 0

    gradient_memory = params * precision_bytes

    total_no_activation = weight_memory + optimizer_memory + gradient_memory

    if hidden_dim and num_layers:
        activation_per_layer = (
            sequence_length * batch_size_per_gpu * hidden_dim * precision_bytes * 4
        )
        activation_memory = activation_per_layer * num_layers
    else:
        activation_memory = params * precision_bytes * 0.5

    if sharding == "fsdp" or sharding == "zero3":
        weight_memory /= num_gpus
        optimizer_memory /= num_gpus
        gradient_memory /= num_gpus
    elif sharding == "zero2":
        optimizer_memory /= num_gpus
        gradient_memory /= num_gpus
    elif sharding == "zero1":
        optimizer_memory /= num_gpus

    per_gpu_total = weight_memory + optimizer_memory + gradient_memory + activation_memory

    return {
        "params_billions": params_billions,
        "weights_gb": weight_memory / 1e9,
        "optimizer_gb": optimizer_memory / 1e9,
        "gradients_gb": gradient_memory / 1e9,
        "activations_gb": activation_memory / 1e9,
        "per_gpu_total_gb": per_gpu_total / 1e9,
        "total_across_gpus_gb": per_gpu_total * num_gpus / 1e9,
        "fits_on_80gb": per_gpu_total / 1e9 <= 80,
        "num_gpus": num_gpus,
        "sharding": sharding,
    }

این ماشین‌حساب به سوالی پاسخ می‌دهد که هر مهندس یادگیری ماشین می‌پرسد: «به چند کارت گرافیک نیاز دارم؟» اندازه‌ی مدل را به آن بدهید و ببینید آیا جا می‌شود یا خیر. استراتژی شاردینگ را آن‌قدر تغییر دهید تا مجموع حافظه به ازای هر کارت گرافیک به زیر 80GB برسد.

گام 5: شبیه‌سازی دقت ترکیبی

مصرف حافظه را بین آموزش‌های FP32، FP16 و دقت ترکیبی مقایسه کنید.

def mixed_precision_comparison(params_billions):
    params = params_billions * 1e9

    fp32_weights = params * 4
    fp32_optimizer = params * 4 * 2
    fp32_gradients = params * 4
    fp32_total = fp32_weights + fp32_optimizer + fp32_gradients

    fp16_weights = params * 2
    fp16_master = params * 4
    fp16_optimizer = params * 4 * 2
    fp16_gradients = params * 2
    fp16_total = fp16_weights + fp16_master + fp16_optimizer + fp16_gradients

    mixed_weights = params * 2
    mixed_optimizer = params * 4 * 2
    mixed_gradients = params * 2
    mixed_total = mixed_weights + mixed_optimizer + mixed_gradients

    return {
        "fp32_total_gb": fp32_total / 1e9,
        "fp16_with_master_gb": fp16_total / 1e9,
        "mixed_bf16_gb": mixed_total / 1e9,
        "savings_vs_fp32": 1 - mixed_total / fp32_total,
    }

بزرگ‌ترین غافلگیری برای بیشتر افراد: دقت ترکیبی مصرف حافظه را نصف نمی‌کند. وضعیت‌های بهینه‌ساز (m و v در Adam) صرف‌نظر از دقت مدل، در فرمت FP32 باقی می‌مانند. برای یک مدل 7B، آموزش FP32 حدود 112GB مصرف می‌کند. دقت ترکیبی 84GB استفاده می‌کند. این یک کاهش 25 درصدی است، نه 50 درصدی؛ زیرا بهینه‌ساز بخش اصلی حافظه را مصرف می‌کند.

استفاده از ابزارهای واقعی

اجرای همه‌ی شبیه‌سازی‌ها

def run_all_demos():
    print("=" * 70)
    print("DATA PARALLELISM SIMULATION")
    print("=" * 70)

    np.random.seed(42)
    data = np.random.randn(64, 32)
    weight = np.random.randn(32, 16)

    def model_fn(batch):
        output = batch @ weight
        loss = np.mean(output ** 2)
        grad = 2 * batch.T @ (batch @ weight) / len(batch)
        return loss, grad

    for n_gpus in [1, 2, 4, 8]:
        loss, grad = simulate_data_parallelism(data, n_gpus, model_fn)
        print(f"  {n_gpus} GPUs: loss={loss:.4f}, grad_norm={np.linalg.norm(grad):.4f}")

    print()
    print("=" * 70)
    print("TENSOR PARALLELISM SIMULATION")
    print("=" * 70)

    x = np.random.randn(4, 8192)
    W = np.random.randn(8192, 8192)

    for n_gpus in [1, 2, 4, 8]:
        output, error = simulate_tensor_parallelism(x, W, n_gpus)
        print(f"  {n_gpus} GPUs: output_shape={output.shape}, max_error={error:.2e}")

    print()
    print("=" * 70)
    print("PIPELINE PARALLELISM SIMULATION")
    print("=" * 70)

    for n_mb in [1, 4, 8, 16, 32]:
        _, total_t, bubble = simulate_pipeline_parallelism(32, 4, n_mb)
        print(f"  {n_mb:2d} micro-batches: total_time={total_t:4d}, bubble={bubble:.1%}")

    print()
    print("=" * 70)
    print("MEMORY CALCULATOR")
    print("=" * 70)

    configs = [
        (7, "none", 1),
        (7, "fsdp", 8),
        (70, "none", 1),
        (70, "fsdp", 8),
        (70, "fsdp", 16),
        (405, "fsdp", 64),
        (405, "fsdp", 128),
    ]

    print(f"  {'Model':>8} {'Sharding':>8} {'GPUs':>5} {'Per-GPU':>10} {'Fits 80GB':>10}")
    print("  " + "-" * 50)
    for params, shard, gpus in configs:
        result = memory_calculator(params, num_gpus=gpus, sharding=shard)
        fits = "Yes" if result["fits_on_80gb"] else "No"
        print(f"  {params:>6}B {shard:>8} {gpus:>5} {result['per_gpu_total_gb']:>8.1f}GB {fits:>10}")

    print()
    print("=" * 70)
    print("MIXED PRECISION COMPARISON")
    print("=" * 70)

    for params_b in [7, 13, 70, 405]:
        result = mixed_precision_comparison(params_b)
        print(f"  {params_b}B: FP32={result['fp32_total_gb']:.0f}GB, "
              f"Mixed BF16={result['mixed_bf16_gb']:.0f}GB, "
              f"Savings={result['savings_vs_fp32']:.0%}")

خروجی این درس

این درس فایل outputs/prompt-distributed-training-planner.md را تولید می‌کند؛ پرامپتی که اندازه‌ی مدل و سخت‌افزار در دسترس را می‌گیرد و یک برنامه‌ی کامل آموزش توزیع‌شده ارائه می‌دهد: استراتژی موازی‌سازی، بودجه‌ی حافظه، سربار ارتباطات و نرخ پردازش مورد انتظار.

تمرین‌ها

1. ماشین‌حساب حافظه را طوری تغییر دهید که شامل چک‌پوینت‌گذاری اکتیویشن (activation checkpointing) شود. با چک‌پوینت‌گذاری، اکتیویشن‌ها را تنها در هر لایه‌ی Kام ذخیره کنید (حالت معمول K=1 است، به معنی محاسبه‌ی دوباره‌ی همه). موازنه‌ی حافظه-محاسبات را نشان دهید: چک‌پوینت‌گذاری چقدر در حافظه صرفه‌جویی می‌کند و چقدر سرعت آموزش را کاهش می‌دهد (تقریباً 33% محاسبات بیشتر برای چک‌پوینت‌گذاری کامل)؟

2. شبیه‌سازی موازی‌سازی پایپ‌لاین را برای پیاده‌سازی زمان‌بندی 1F1B (یک جلو، یک عقب) که در PipeDream استفاده می‌شود، گسترش دهید. کسر حباب را در مقایسه با زمان‌بندی ساده برای 4 مرحله و 8 ریزدسته بسنجید. زمان‌بندی 1F1B باید اوج مصرف حافظه‌ی (peak memory) کمتری داشته باشد زیرا گام‌های رو به عقب را زودتر آغاز می‌کند.

3. یک شبیه‌ساز انباشت گرادیان (Gradient Accumulation) پیاده کنید. به جای all-reduce بعد از هر ریزدسته، گرادیان‌ها را به صورت محلی برای K گام انباشت کنید و سپس all-reduce را اجرا نمایید. نشان دهید که چگونه این کار ارتباطات را K برابر کاهش می‌دهد ولی گرادیان‌های نهایی یکسانی (و در نتیجه آموزش یکسانی) تولید می‌کند.

4. یک برآوردگر هزینه بسازید. با دریافت اندازه‌ی مدل، تعداد توکن‌های هدف، نوع کارت گرافیک (A100 ساعتی 2 دلار، H100 ساعتی 3.50 دلار) و استراتژی موازی‌سازی، هزینه‌ی کل آموزش را به دلار تخمین بزنید. آن را با هزینه‌های شناخته‌شده اعتبارسنجی کنید: طبق گزارش‌ها Llama 3 405B حدود 100 میلیون دلار و DeepSeek V3 حدود 5.6 میلیون دلار هزینه داشته‌اند.

5. قابلیت ZeRO-Offload را به ماشین‌حساب حافظه اضافه کنید. فرض کنید رم CPU برابر 512GB به ازای هر نود و حافظه‌ی NVMe برابر 2TB است. نشان دهید چگونه انتقال وضعیت‌های بهینه‌ساز به CPU اجازه می‌دهد یک مدل 70B به جای 16 کارت، روی 4 کارت گرافیک آموزش داده شود؛ هرچند به قیمت 30 تا 50 درصد کندتر شدن گام‌های بهینه‌ساز تمام می‌شود.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• Rajbhandari et al., 2020 -- "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models" -- مقاله‌ی DeepSpeed ZeRO که مراحل سه‌گانه‌ی شاردینگ را تعریف کرد
• Shoeybi et al., 2020 -- "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism" -- موازی‌سازی تنسور ارائه‌شده توسط شرکت NVIDIA برای ترنسفورمرها
• Narayanan et al., 2021 -- "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM" -- موازی‌سازی سه‌بعدی با ترکیب روش‌های داده، تنسور و پایپ‌لاین
• Zhao et al., 2023 -- "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel" -- پیاده‌سازی بومی FSDP در PyTorch
• گزارش فنی Llama 3 -- آموزش روی 16,384 کارت گرافیک با جزئیات موازی‌سازی سه‌بعدی
• گزارش فنی DeepSeek-V3 -- چگونگی کاهش هزینه‌ی آموزش به میزان یک مرتبه بزرگی با معماری MoE