انباشت گرادیان (Gradient Accumulation): آموزش دسته‌های بزرگ با حافظه محدود

یکی از موانع همیشگی در آموزش مدل‌های یادگیری عمیق، محدودیت حافظه‌ی کارت‌های گرافیک (GPU VRAM) است. در شرایطی که تهیه یا اجاره‌ی GPUهای گران‌قیمت خارجی دشوار و هزینه‌بر است، بسیاری از توسعه‌دهندگان ناچارند فرایند یادگیری را روی کارت‌های گرافیک شخصی یا منابع سخت‌افزاری محدود پیش ببرند. در این حالت، انتخاب دسته‌های داده‌ی بزرگ (Large Batch Size) که باعث پایداری گرادیان و هموار شدن فرایند همگرایی می‌شوند، خیلی زود به خطای کمبود حافظه (CUDA Out of Memory) می‌انجامد.

تکنیک انباشت گرادیان (Gradient Accumulation) راهکار استاندارد مهندسی برای دور زدن این بن‌بست سخت‌افزاری است. با این روش، شما داده‌های بزرگ را به زیربخش‌های کوچک‌تر (میکروبچ) تقسیم می‌کنید، گرادیان هر بخش را بدون به‌روزرسانی زودهنگام پارامترها جمع می‌زنید و در نهایت یک گام بهینه‌سازیِ یکنواخت و معادل با یک دسته‌ی کامل برمی‌دارید. این درس منطق ریاضی، نحوه مقیاس‌بندی صحیح تابع زیان و پیاده‌سازی گام‌به‌گام این تکنیک را در PyTorch آموزش می‌دهد.

با دسته‌ای مؤثر (effective batch) که حافظه‌تان توان پردازش یک‌باره‌اش را ندارد آموزش دهید؛ میکروبچ به میکروبچ. تابع زیان را مقیاس‌بندی کنید، گام بهینه‌ساز را نگه دارید و بگذارید گرادیان‌ها روی هم انباشته شوند.

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

• به‌دست‌آوردن رابطه‌ی دسته‌ی مؤثر: effective_batch = micro_batch * accum_steps * world_size که روی تک‌دستگاهِ این درس به صورت micro_batch * accum_steps درمی‌آید.
• پیاده‌سازی مقیاس‌بندی تابع زیان به ازای هر میکروبچ (micro-batch) به‌گونه‌ای که گرادیانِ انباشته‌شده دقیقاً با گام رو به عقبِ (backward pass) یک دسته‌ی کامل همخوانی داشته باشد.
• صرف‌نظر کردن از همگام‌سازی بهینه‌ساز تا رسیدن به آخرین میکروبچ (sync-on-last-step).
• خواندن نمودار نرخ پردازش (throughput) برحسب دسته‌ی مؤثر و تحلیل بازده نزولی آن.

صورت مسئله

می‌خواهید مدل را با دسته‌ی مؤثر 512 آموزش دهید؛ چون در این مقیاس، منحنی تابع زیان (loss) هموارتر است و گام بهینه‌ساز (optimizer step) معنای بهتری پیدا می‌کند. اما شتاب‌دهنده‌ی روی میزکارتان بعد از 32 نمونه با خطای کمبود حافظه متوقف می‌شود. دو برابر کردن دسته ممکن نیست؛ نصف کردن مدل هم راه به جایی نمی‌برد. ترفندی که این حوزه در سال 2017 سراغش رفت و دیگر کنارش نگذاشت، اجرای 16 گام رو به عقب، انباشتن گرادیان‌ها درون بافرهای پارامتر، و اجرای گام بهینه‌ساز تنها پس از رسیدن شمارنده به هدفِ تعیین‌شده است.

ریسک کار اینجاست: تابع زیان دیگر همان عددی نیست که با دسته‌ی بزرگ‌تر به دست می‌آمد. اگر انتروپی متقاطع (cross entropy)ِ 16 دسته‌ی کوچک (mini-batch) را ساده با هم جمع کنید، 16 برابرِ تابع زیانِ یک دسته‌ی کامل می‌شود. بدون مقیاس‌بندی، جهت گرادیان درست می‌ماند اما مقدارش اشتباه است و گام بهینه‌ساز 16 برابر بیش از اندازه برداشته می‌شود. راه‌حل فقط یک تقسیم ساده است؛ تقسیمی که فراموش کردنش هم بسیار ساده است.

مفاهیم پایه

قرارداد کار کوتاه و مشخص است:

• پیش از فراخوانی backward()، تابع زیانِ هر میکروبچ بر accum_steps تقسیم می‌شود. PyTorch به طور پیش‌فرض گرادیان‌ها را درون param.grad جمع می‌زند؛ این تقسیمْ مجموعِ در حال محاسبه را به مقیاس درست بازمی‌گرداند.
• گام بهینه‌ساز فقط یک بار به ازای هر دسته‌ی مؤثر اجرا می‌شود؛ یعنی پس از گام رو به عقبِ آخرین میکروبچ. اگر وسط عملیات انباشت گامی بردارید، همه‌ی پارامترهایی که ادامه‌ی اجرا به آن‌ها وابسته است دچار انحراف می‌شوند.
• وضعیت درونی بهینه‌ساز (بافرهای تکانه [momentum] و گشتاورهای Adam) به ازای هر گام مؤثر جلو می‌رود، نه به ازای هر میکروبچ. در غیر این صورت، میانگین‌های متحرک نمایی بسامد اشتباهی دریافت می‌کنند و برنامه‌ی زمان‌بندی نرخ یادگیری (learning rate schedule) خیلی زود می‌سوزد و هدر می‌رود.
• روی یک دستگاه، این کار صرفاً حساب‌وکتاب داخلی است. در خوشه‌های چندرتبه‌ای (multi-rank)، همین الگو میکروبچ‌های غیرپایانی را درون زمینه‌ی no_sync قرار می‌دهد تا عملیات all-reduce روی گرادیان‌ها اجرا نشود؛ به این ترتیب آخرین میکروبچ، کل گرادیان انباشته‌شده را در یک گذر کاهش می‌دهد و به‌جای پرداخت N بار هزینه‌ی شبکه، این هزینه فقط یک بار پرداخت می‌شود.

اثبات معادل بودن در قالب کد

loss = criterion(model(x_full), y_full)
loss.backward()
opt.step()

معادل است با:

for x, y in chunks(x_full, y_full, n):
    scaled = criterion(model(x), y) / n
    scaled.backward()
opt.step()

تنها تا حد ترتیب جمع‌زدن اعداد ممیز شناور تفاوت دارد. بافر گرادیانِ انباشته‌شده در پایان حلقه، دقیقاً همان تنسوری (tensor) است که گام رو به عقب روی یک دسته‌ی کامل تولید می‌کرد. کد این درس، این هم‌ارزی را با بررسی حداکثر اختلاف مطلق زیر 1e-4 در equivalence_check تضمین می‌کند.

هزینه صرف چه چیزی می‌شود؟

هر میکروبچ به یک گذر رو به جلو (forward pass) و یک گام رو به عقب نیاز دارد. با انباشت، شما حافظه را با زمان مبادله می‌کنید. منحنی نرخ پردازش در فایل outputs/accum-curve.json نشان می‌دهد با ثابت نگه‌داشتن میکروبچ و بزرگ‌تر کردن دسته‌ی مؤثر، چه اتفاقی رخ می‌دهد:

ناهار مجانی در کار نیست. دو برابر کردن accum_steps، زمان واقعیِ (wall time) هر گام بهینه‌ساز را هم دو برابر می‌کند. چیزی که دستخوش تغییر می‌شود، واریانس تخمین گرادیان است: با همان بودجه‌ی زمانیِ سپری‌شده، گام‌های بهینه‌ساز کمتری برداشته‌اید، اما هر کدام روی نمونه‌های بیشتری میانگین گرفته شده است. مقالات علمی دسته‌ی بزرگ و دسته‌ی کوچک را دو مسئله‌ی بهینه‌سازی متفاوت می‌دانند؛ درسِ اینجا مکانیکی است، نه آماری.

پیاده‌سازی

فایل code/main.py دست‌سازه‌ی (artifact) اجرایی است و سه کار انجام می‌دهد:

گام 1: بررسی هم‌ارزی

تابع equivalence_check() دو نسخه‌ی یکسان از یک شبکه را با سید (seed) مشابه می‌سازد. نسخه‌ی اول یک دسته‌ی 16 نمونه‌ای را در یک گذر رو به جلو پردازش می‌کند. نسخه‌ی دوم چهار بخشِ 4 نمونه‌ای را می‌بیند که تابع زیان آن‌ها بر چهار تقسیم شده است. این تابع، بافرهای گرادیان را پیش از گام بهینه‌ساز و پارامترها را پس از آن مقایسه می‌کند. شرط تایید این است: max_abs_diff < 1e-4.

گام 2: الگوی همگام‌سازی در گام پایانی (sync-on-last-step)

تابع train_one_optimizer_step روی میکروبچ‌ها پیمایش می‌کند. برای همه‌ی میکروبچ‌ها به جز آخری، وارد no_sync_context(model) می‌شود. در یک فرایندِ (process) تک‌دستگاهی، این زمینه هیچ کاری نمی‌کند (no-op)؛ اما در DDP، درست همان‌جایی است که all-reduce گرادیان نادیده گرفته می‌شود. حساب‌وکتاب در هر دو حالت یکسان است. یک sync_counter دفعات خروج از دامنه‌ی (scope) no_sync را ثبت می‌کند؛ برای N میکروبچ، این شمارشگر به ازای هر گام مؤثر عدد یک را نشان می‌دهد، نه N.

گام 3: منحنی نرخ پردازش

تابع sweep_effective_batches همان مدل را با یک میکروبچ ثابت و فهرستی از گام‌های انباشت مختلف اجرا می‌کند. برای هر حالت، معیارهای زیر ثبت می‌شوند:

• samples_per_sec: کل نمونه‌های دیده‌شده تقسیم بر زمان واقعی
• median_step_ms: صدک پنجاهم زمان به ازای هر گام مؤثر
• sync_calls: نقاط ارتباط جمعیِ اجراشده
• avg_loss: میانگین تابع زیان در گام‌های بهینه‌ساز در طول این پیمایش

خروجی در فایل outputs/accum-curve.json ذخیره می‌شود و در نوت‌بوک‌ها قابل استفاده است.

آن را اجرا کنید:

python3 code/main.py

اسکریپت ابتدا اختلاف هم‌ارزی، سپس جدول پیمایش و در نهایت مسیر فایل JSON را چاپ می‌کند. کد خروجی نیز صفر خواهد بود.

استفاده در عمل

در آموزش محیط عملیاتی (پروداکشن)، انباشت گرادیان پشت یک دکمه‌ی تنظیم قرار دارد. الگوی مرسوم در PyTorch به شکل accumulation_steps = effective_batch // (micro_batch * world_size) تعریف می‌شود. فریم‌ورک‌هایی که در این درس اجازه‌ی استفاده از آن‌ها را ندارید هم دقیقاً همین حلقه را در خود پیچیده‌اند و گام‌ها یکسان است: مقیاس‌بندی تابع زیان، رد کردن همگام‌سازی در میکروبچ‌های غیرپایانی، انباشت و در نهایت یک بار اجرای گام بهینه‌ساز.

سه الگو در پروژه‌های واقعی:

• اندازه‌ی میکروبچ طوری انتخاب می‌شود که حافظه‌ی دستگاه تا مرز اشباع پر شود. هر مقداری کمتر از آن، چرخه‌های پردازشی شتاب‌دهنده را هدر می‌دهد؛ هر مقداری بیشتر از آن، برنامه را متوقف می‌کند.
• دسته‌ی مؤثر بر اساس برنامه‌ی زمان‌بندی نرخ یادگیری انتخاب می‌شود. دسته‌های مؤثر بزرگ نیازمند وارم‌آپ (پیش‌گرمایش) و نرخ‌های یادگیریِ مقیاس‌بندی‌شده هستند؛ این همان قاعده‌ی مقیاس‌بندی خطی است که از سال 2017 درباره‌اش صحبت می‌شود.
• تعداد دفعات انباشت، رابطِ (bridge) میان این دو است و تنها تنظیمی است که در محیط اجرا (runtime) می‌توانید بدون بازنویسی دیتا لودر تغییرش دهید.

تحویل خروجی

فایل outputs/skill-gradient-accumulation.md این دستورالعمل را مستند کرده تا همکار دیگری بتواند آن را در یک ریپوی جدید به کار بگیرد: تقسیم تابع زیان بر accum_steps، صرف‌نظر از همگام‌سازی بهینه‌ساز در میکروبچ‌های غیرپایانی، اجرای گام بهینه‌ساز یک بار به ازای هر دسته‌ی مؤثر، و ثبت لاگ نرخ پردازش در برابر دسته‌ی مؤثر در قالب یک فایل JSON تا توازن میان این دو به وضوح دیده شود.

تمرین‌ها

1. پیمایش را با پارامتر --num-steps 100 دوباره اجرا کنید و نمودار نمونه بر ثانیه را بر حسب دسته‌ی مؤثر رسم کنید. این منحنی در کجا افقی می‌شود؟
2. یک حالت مقیاس‌بندی اشتباه اضافه کنید (بدون عمل تقسیم) و تفاوت پارامترها را در گام 1 نسبت به حالت مرجع نشان دهید.
3. الگوریتم SGD را با AdamW جایگزین کنید و تأیید کنید که وضعیت بهینه‌ساز فقط یک بار به ازای هر گام مؤثر جلو می‌رود، نه در هر میکروبچ.
4. یک بسته‌بند واقعیِ DistributedDataParallel پیاده کنید و no_sync_context را به متد آن متصل کنید. بررسی کنید که به ازای هر دسته‌ی مؤثر، تعداد sync_calls دقیقاً N-1 بار کاهش یابد.
5. بررسی هم‌ارزی را طوری تغییر دهید که دو شیوه‌ی مختلف تقسیم میکروبچ (2 دسته‌ی 8تایی در برابر 4 دسته‌ی 4تایی) را مقایسه کند و توضیح دهید چرا باید آستانه‌ی خطا (tolerance) را کمی بازتر بگذارید.

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

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

• مستندات PyTorch درباره‌ی DistributedDataParallel.no_sync برای دیدن پیاده‌سازی پروداکشنِ ترفند sync-on-last-step.
• مقاله‌ی گویال و همکاران (Goyal et al., 2017) در مورد مقیاس‌بندی خطی برای آموزش با دسته‌های بزرگ؛ مرجع اصلی برای پی‌بردن به اهمیت دسته‌ی مؤثر.
• پیگیری باگ‌ها (Issue tracker) در PyTorch پیرامون تداخل‌های انباشت گرادیان با نامقیاس‌بندی در دقت ترکیبی (mixed precision).
• درس‌های 42 تا 45 از فاز 19 که داربستِ (scaffolding) مربوط به مدل، دیتا لودر، بهینه‌ساز و فرایند آموزش را که این درس فرض گرفته پوشش می‌دهند.
• درس 47 از فاز 19 که به چک‌پوینت (checkpoint) و ادامه‌ی آموزش می‌پردازد تا اجرای انباشت‌های طولانی در برابر محدودیت‌های زمانی محیط متوقف نشود.