خرد کردن وضعیت بهینه‌ساز با ZeRO-1: پیاده‌سازی و تحلیل حافظه

آموزش مدل‌های زبانی چندمیلیارد پارامتری چالش بزرگی به نام کمبود حافظه پردازنده گرافیکی (VRAM) به همراه دارد. در آموزش به روش موازی‌سازی داده استاندارد (DDP)، همه اطلاعات شامل وزن‌ها، گرادیان‌ها و وضعیت بهینه‌ساز روی تک‌تک کارت‌ها کپی می‌شوند. وضعیت بهینه‌ساز Adam به‌تنهایی بیش از نیمی از فضای حافظه را می‌بلعد و به مانع اصلی در آموزش مدل‌های بزرگ تبدیل می‌شود.

با توجه به اینکه تهیه و اجاره GPUهای مدرن با حافظه بالا برای توسعه‌دهندگان ایرانی هزینه‌بر و دسترسی به آن‌ها دشوار است، بهینه‌سازی مصرف حافظه اهمیت دوچندانی پیدا می‌کند. با روش‌هایی مثل خرد کردن وضعیت بهینه‌ساز (ZeRO Stage 1)، می‌توان همان مدل‌ها را روی سخت‌افزارهای دردسترس‌تر با ظرفیت حافظه محدودتر به پیش برد.

در این درس یاد می‌گیرید که چگونه با تکنیک ZeRO-1 وضعیت الگوریتم Adam و کپی با دقت float32 را میان کارت‌های مختلف تقسیم کنید، با عملیات reduce_scatter و allgather ترافیک ارتباطی را بدون تغییر نگه دارید و بدون افت سرعت، مصرف حافظه را به شکل خطی کاهش دهید.

الگوریتم Adam برای هر پارامتر دو تخمین گشتاور (moment estimates) ذخیره می‌کند که هر دو از نوع float32 هستند. یک مدل 7 میلیارد پارامتری (7B) حدود 56 گیگابایت وضعیت بهینه‌ساز (optimiser state) به همراه دارد. الگوریتم ZeRO stage 1 این حافظه را بین N رتبه (rank) خرد (shard) می‌کند؛ یعنی هر رتبه فقط مالک 1/N از بهینه‌ساز خواهد بود. پس از گام بهینه‌سازی محلی، قطعه‌های به‌روزشده‌ی پارامترها بازپخش (broadcast) می‌شوند، همه‌ی رتبه‌ها مدل کامل را بازسازی می‌کنند و گام بعدی آغاز می‌شود. دستاورد این کار، کاهش خطی مصرف حافظه در بزرگ‌ترین بخش تخصیص‌یافته در کل استک آموزش است.

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

• خرد کردن وضعیت بهینه‌ساز (گشتاور اول، گشتاور دوم و نسخه‌ی اصلی با دقت fp32) میان N رتبه به‌طوری که سهم هر رتبه 1/N باشد.
• استفاده از reduce_scatter برای رساندن مجموع گرادیان‌های مربوط به هر قطعه فقط به رتبه‌ی مالک آن، و سپس بهره‌گیری از allgather برای بازپخش مجدد قطعه‌های به‌روزشده‌ی پارامترها.
• محاسبه‌ی جدول صرفه‌جویی حافظه برای ZeRO-1، ZeRO-2 و ZeRO-3 در مقایسه با DDP معمولی.
• توجیه انتخاب بین ZeRO-1، ZeRO-2 یا ZeRO-3 بر اساس ابعاد مدل و بودجه‌ی پهنای باند شبکه.

صورت مسئله

روش معمولی DDP همه‌چیز را تکثیر می‌کند: پارامترها، گرادیان‌ها و وضعیت بهینه‌ساز به‌طور کامل روی تک‌تک رتبه‌ها وجود دارند. برای یک مدل 7B با دقت fp16، این موضوع به معنای 14 گیگابایت پارامتر، 14 گیگابایت گرادیان و 56 گیگابایت وضعیت بهینه‌ساز (دو گشتاور float32 الگوریتم Adam) روی هر رتبه است. علاوه‌بر این، آموزش با دقت ترکیبی (Mixed Precision) یک نسخه‌ی اصلی (master copy) از وزن‌ها با دقت float32 به حجم 28 گیگابایت نیز به این ارقام اضافه می‌کند. وضعیت بهینه‌ساز بزرگ‌ترین بخش حافظه را اشغال می‌کند و در عین حال ساده‌ترین بخش برای خرد کردن است؛ زیرا فقط در طول گام به‌روزرسانی پارامترها با آن سروکار داریم، نه در مرحله‌ی رفت (forward) یا برگشت (backward).

روش ZeRO stage 1 وضعیت بهینه‌ساز را خرد می‌کند. هر رتبه 1/N از گشتاورهای Adam را نگه می‌دارد. پس از مرحله‌ی برگشت، به‌جای اجرای allreduce روی کل گرادیان و اِعمال گام به‌روزرسانی به‌صورت محلی، ZeRO از reduce_scatter استفاده می‌کند تا هر رتبه فقط گرادیان‌های جمع‌شده‌ی مربوط به قطعه‌ی خودش را تحویل بگیرد. سپس رتبه، گام بهینه‌ساز را تنها روی قطعه‌ی مربوط به پارامترهای اصلی خود اعمال می‌کند. قطعه‌های به‌روزشده‌ی پارامترها در ادامه از طریق allgather بازپخش می‌شوند تا همه‌ی رتبه‌ها مدل کامل را برای مرحله‌ی رفت بعدی در اختیار داشته باشند. حافظه‌ی مصرفی بهینه‌ساز بر N تقسیم می‌شود. ترافیک شبکه در هر گام دقیقاً معادل DDP است: از نظر پهنای باند مصرفی، یک reduce_scatter به‌اضافه‌ی یک allgather برابر است با یک allreduce. نتیجه این‌که حافظه آزاد می‌شود و نرخ پردازش (throughput) ثابت می‌ماند.

مفاهیم پایه

مراحل مختلف ZeRO

مرحله‌ی اول (Stage 1) کم‌هزینه‌ترین برد را به همراه دارد، زیرا وضعیت بهینه‌ساز بیشترین سهم را در بودجه‌ی حافظه دارد. مرحله‌ی دوم (Stage 2) نیازمند منطق تجمیع قطعه‌های گرادیان است، اما پهنای باند آن تفاوتی نمی‌کند. مرحله‌ی سوم (Stage 3 یا همان FSDP) به ازای هر مرحله‌ی رفت و برگشت هزینه‌ی ارتباطی لایه‌به‌لایه را متحمل می‌شود تا در عوض حافظه‌ی مربوط به قطعه‌های پارامترها را هم آزاد کند. در این درس مرحله‌ی اول را به‌طور کامل پیاده‌سازی می‌کنیم.

محاسبات ریاضی حافظه با اعداد واقعی

برای مدلی با P پارامتر که با الگوریتم Adam و دقت ترکیبی آموزش می‌بیند:

در حالت N=8: روش معمولی 16P بایت و ZeRO-1 برابر 5.5P بایت است که 65% کاهش نشان می‌دهد. در حالت N=64: روش معمولی 16P بایت و ZeRO-1 برابر 4.19P بایت است که 74% کاهش به همراه دارد.

چرا reduce_scatter از ترکیب allreduce-سپس-خردکردن بهتر است؟

عملیات Allreduce گرادیان جمع‌شده‌ی کامل را در اختیار همه‌ی رتبه‌ها می‌گذارد. اگر تنها به قطعه‌ی r نیاز داشته باشید، عملاً بخش (N-1)/N از گرادیان جمع‌شده روی رتبه‌ی r هدر می‌رود. عملیات Reduce_scatter دقیقاً همان قطعه‌ای را تحویل می‌دهد که رتبه مالک آن است؛ حجم بایت‌های ردوبدل‌شده به ازای هر رتبه معادل allreduce است (چون allreduce ترکیبی از reduce_scatter به‌اضافه‌ی allgather است)، اما نیمه‌ی دوم بعداً با allgather مربوط به قطعه‌های پارامترها جایگزین می‌شود. ترافیک خالص شبکه کاملاً یکسان با DDP است، در حالی که حافظه بر N تقسیم می‌شود.

پیاده‌سازی

فایل code/main.py موارد زیر را پیاده‌سازی می‌کند:

• توابع flatten_params(module) و unflatten_into(module, flat) که پارامترهای مدل را داخل یک تنسور پیوسته در حافظه (contiguous) بسته‌بندی کرده و سپس به حالت اولیه بازمی‌گردانند. ساختار مسطح و پیوسته همان چیزی است که خرد کردن بر اساس رتبه را به یک برش (slice) ساده تبدیل می‌کند.
• کلاس ZeroOptimizer(model, world_size, rank, lr) که مالک قطعه‌ی مربوط به رتبه از نسخه‌ی اصلی و گشتاورهای Adam است.
• متد step() که reduce_scatter را روی گرادیان مسطح اجرا می‌کند، الگوریتم Adam را روی قطعه‌ی همان رتبه اعمال کرده و پارامترهای به‌روزشده را با allgather بازپخش می‌کند.
• یک نمونه‌ی نمایشی که یک پرسپترون چندلایه (MLP) سه‌لایه را برای 20 گام آموزش می‌دهد و بودجه‌ی حافظه‌ی مصرفی در هر گام را در کنار یک خط مبنا (baseline) از DDP معمولی چاپ می‌کند.

اجرای برنامه:

python3 code/main.py

خروجی: تابع زیان (loss) در هر گام و جدول حافظه‌ای که نشان می‌دهد ZeRO-1 برخلاف نسخه‌ی کامل DDP، تنها 1/N از وضعیت بهینه‌ساز را روی هر رتبه نگه می‌دارد.

الگوهای کاربردی در محیط عملیاتی

سه الگو برای آماده‌سازی ZeRO جهت استقرار در محیط عملیاتی (پروداکشن) ضروری است:

چک‌پوینت‌گیری خردشده (Sharded Checkpointing) اهمیت حیاتی دارد. وضعیت بهینه‌ساز در ZeRO-1 میان رتبه‌ها تقسیم شده است؛ بنابراین چک‌پوینت باید ثبت کند که چه رتبه‌ای مالک چه بخشی است. در درس 80 مانیفست (manifest) چک‌پوینت خردشده را می‌سازیم که امکان ادامه‌ی اجرای ZeRO را با همان تعداد رتبه (world size) فراهم می‌کند. بدون این مانیفست، وضعیت ذخیره‌شده هنگام راه‌اندازی مجدد غیرقابل خواندن خواهد بود.

اصل ماجرا دقت ترکیبی است. روش ZeRO یک تکنیک مبتنی بر دقت ترکیبی است؛ چیزی که خرد می‌شود همان نسخه‌ی اصلی با دقت fp32 است. اجرای ZeRO بدون دقت ترکیبی، هزینه‌ی حافظه‌ی نسخه‌ی اصلی fp32 را به سیستم تحمیل می‌کند بدون اینکه صرفه‌جویی حاصل از دقت fp16 در مرحله‌ی رفت را در بر داشته باشد. اجراهای محیط عملیاتی همیشه ZeRO را با autocast یا وزن‌های bf16 همراه می‌کنند.

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

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

الگوهای محیط عملیاتی:

• DeepSpeed ZeRO: پیاده‌سازی مرجع. فایل deepspeed_config.json مراحل 1، 2 یا 3 و ابعاد پارتیشن‌ها را تعیین می‌کند.
• PyTorch FSDP: معادل بومی پایتورچ. مقدار ShardingStrategy.SHARD_GRAD_OP همان ZeRO-2 است و FULL_SHARD معادل ZeRO-3 خواهد بود.
• HuggingFace Accelerate: هر دو ابزار DeepSpeed و FSDP را ذیل یک پیکربندی یکپارچه ارائه می‌دهد.

جمع‌بندی و ادامه‌ی مسیر

درس 79 (موازی‌سازی پایپ‌لاینی یا Pipeline Parallel) بُعد عمود بر این خردسازی را بررسی می‌کند: به‌جای خرد کردن وضعیت بهینه‌سازِ یک مدل واحد، پایپ‌لاین لایه‌ها را میان رتبه‌ها تقسیم می‌کند. در درس 81 ترکیب DDP و ZeRO را در قالب یک دموی سرتاسری (end-to-end) پیاده‌سازی می‌کنیم.

تمرین‌ها

1. کد را با خرد کردن گرادیان‌ها به ZeRO-2 ارتقا دهید: هر رتبه فقط باید گرادیان مربوط به قطعه‌ی خودش را ذخیره کند؛ این کار با صفر کردن بخش‌های غیرمربوط بعد از مرحله‌ی برگشت انجام می‌شود.
2. یک پروفایلر حافظه اضافه کنید که مصرف واقعی بایت‌های fp32 را روی رتبه‌ی 0 چاپ کرده و با مقدار پیش‌بینی‌شده در فرمول مقایسه کند.
3. زمان اجرای واقعی (wall-clock time) در هر گام را بین DDP معمولی و ZeRO-1 اندازه بگیرید و آن را به بخش‌های رفت، برگشت و ارتباطات شبکه تفکیک کنید.
4. قابلیت برش گرادیان (Gradient clipping) را تحت ZeRO-1 پیاده‌سازی کنید: نُرم L2 باید روی همه‌ی قطعه‌ها از طریق allreduce گرفتن از توان دوم نرم محلی محاسبه شود.
5. یک «ZeRO ساده‌لوحانه» (naive ZeRO) با allreduce به‌جای reduce_scatter پیاده کنید و تفاوت زمان تبادل شبکه را بسنجید. با تکیه بر اعداد و ارقام از انتخاب reduce_scatter دفاع کنید.

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

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

• مقاله Rajbhandari و همکاران، ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• مستندات DeepSpeed ZeRO
• مستندات PyTorch FSDP
• فاز 19، درس 76 – بررسی مبانی reduce_scatter و allgather که این درس بر پایه‌ی آن‌ها بنا شده است
• فاز 19، درس 80 – چک‌پوینت‌گیری خردشده که وضعیت ZeRO ملزم به استفاده از آن است