برش گرادیان و آموزش با دقت ترکیبی (AMP) در PyTorch

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

با توجه به اینکه تهیه و اجاره‌ی کارت‌های گرافیک قدرتمند در ایران با چالش‌های پرداخت ارزی و هزینه‌های بالا همراه است، بهینه‌سازی حداکثری توان سخت‌افزاری اهمیت دوچندان پیدا می‌کند. استفاده از آموزش با دقت ترکیبی (Mixed Precision) مصرف حافظه را به نصف می‌رساند و سرعت محاسبات را 2 تا 3 برابر می‌کند؛ اما به دلیل دامنه‌ی محدود محاسبات در قالب FP16، ریسک سرریز گرادیان افزایش می‌یابد.

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

بهینه‌ساز (optimizer) و زمان‌بند (schedule) درس قبل فرض می‌کردند گرادیان‌ها همیشه مقادیر معقولی دارند؛ اما معمولاً این‌طور نیست. یک بچ (batch) داده‌ی نامناسب به‌تنهایی می‌تواند نرم گرادیان را تا سه مرتبه بزرگی (1,000 برابر) بالا ببرد. آموزش با دقت ترکیبی (Mixed Precision) این مشکل را حادتر هم می‌کند، چرا که سرریز (overflow) محاسبات FP16 را به بخش محاسبه‌ی تابع زیان می‌کشاند. در این درس، دو کمربند ایمنی ضروری را می‌سازیم که بدون آن‌ها هیچ فرآیند آموزشی نباید راهی محیط عملیاتی (پروداکشن) شود: برش گرادیان (Gradient clipping) بر اساس یک نرم L2 سراسری مشخص، و یک حلقه‌ی دقت ترکیبی با autocast و GradScaler که مقادیر NaN و Inf را تشخیص دهد، گام بهینه‌ساز را بدون خرابکاری رد کند (skip) و ضریب مقیاس‌بندی را برای خطایابی و تحلیل علت ریشه‌ای ثبت کند.

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

• محاسبه‌ی نرم L2 سراسری روی گرادیان تمام پارامترها و برش درجا (in-place) در صورت فراتر رفتن از آستانه (threshold) تنظیمی.
• قرار دادن گام آموزش در autocast همراه با GradScaler تا گذر رو به جلو و عقب در FP16 از سرریز جان سالم به در ببرند.
• تشخیص NaN و Inf در تابع زیان یا گرادیان، رد کردن گام بهینه‌ساز و ثبت لاگ این رد شدن.
• گزارش ضریب مقیاس‌بندی GradScaler در هر گام تا توالی‌های طولانی رد شدن گام‌ها فوراً مشخص شوند.

صورت مسئله

اجرای فرآیند آموزشی که دیروز کاملاً بی‌نقص پیش می‌رفت، در گام 8,217 ناگهان با نمودار تابع زیانی (loss) مواجه می‌شود که خطی عمودی رو به بالا رسم می‌کند. مقصر ماجرا یک بچ تکی است که نرم گرادیان آن به 4,200 رسیده؛ یعنی 20 برابرِ قله‌ی قبلی. بدون برش گرادیان، بهینه‌ساز گامی برمی‌دارد که تمام یادگیری مدل در یک ساعت گذشته را نابود می‌کند. اما با یک برش L2 سراسری روی نرم 1.0، همان بچ آپدیتی با نرم واحد اعمال می‌کند؛ تابع زیان روی مسیر اصلی خود باقی می‌ماند و فرآیند آموزش نجات پیدا می‌کند.

آموزش با دقت ترکیبی با انجام گذر رو به جلو (forward pass) و بخش عمده‌ی گام رو به عقب (backward pass) در قالب FP16، نرخ پردازش (throughput) را 2 تا 3 برابر افزایش می‌دهد. هزینه‌ی این کار، بازه‌ی توان بسیار محدود در FP16 است. یک گرادیان معمولی که در FP16 سرریز می‌کند به Inf تبدیل می‌شود، سپس در لایه‌های بعدی به‌صورت NaN منتشر شده و در گام بعدی بهینه‌ساز، تمام وزن‌ها را برابر با NaN قرار می‌دهد. ابزار GradScaler در PyTorch این مشکل را حل می‌کند: تابع زیان را قبل از گام رو به عقب در یک ضریب مقیاس‌بندی بزرگ ضرب کرده و گرادیان‌ها را قبل از گام بهینه‌ساز بر همان ضریب تقسیم می‌کند. اگر هنگام بازگرداندن مقیاس (unscale)، گرادیانی برابر با Inf یا NaN باشد، مقیاس‌بند آن گام را رد کرده و ضریب مقیاس‌بندی را نصف می‌کند؛ اگر N گام قبلی بدون خطا سپری شده باشند، ضریب را دو برابر می‌کند. در طول دوره‌ی آموزش، این ضریب بیشترین مقداری را که بازه‌ی FP16 اجازه می‌دهد پیدا خواهد کرد.

چالش اصلی در پیاده‌سازی، اتصال درست این دو بخش به یکدیگر است. اگر برش را قبل از بازگرداندن مقیاس انجام دهید، آستانه روی گرادیان‌های مقیاس‌بندی‌شده اعمال می‌شود؛ اگر برش را بعد از بازگرداندن مقیاس بگذارید، ترتیب عملیات روی GradScaler اهمیت حیاتی پیدا می‌کند. ترتیب صحیح به این شکل است: ابتدا scaler.scale(loss).backward()، سپس scaler.unscale_(optimizer)، بعد clip_grad_norm_، پس از آن scaler.step(optimizer) و در نهایت scaler.update(). هر ترتیب دیگری حلقه‌ای تحویل می‌دهد که در سکوت و بدون خطا، کارکردش مختل شده است.

مفاهیم پایه

نرم L2 سراسری

نرم L2 سراسری در واقع نرم اقلیدسی بردار به‌هم‌پیوسته‌ی گرادیان‌ها است، نه نرم جداگانه برای هر پارامتر. PyTorch این قابلیت را با تابع torch.nn.utils.clip_grad_norm_(parameters, max_norm) پیاده‌سازی کرده است. این تابع نرم قبل از برش را برمی‌گرداند؛ بنابراین می‌توانیم هم مقدار طبیعی و هم مقدار برش‌خورده را لاگ کنیم. ثبت هر دو مقدار برای تشخیص وضعیت «برش مداوم در تمام گام‌ها» کاملاً ضروری است.

ابزارهای autocast و GradScaler

دستور torch.amp.autocast(device_type) یک مدیر بافتار (context manager) است که عملیات‌های واجد شرایط (عمدتاً دسته‌ی ضرب ماتریسی) را به انتخاب خود در قالب FP16 اجرا می‌کند. کلاس torch.amp.GradScaler(device_type) ابزار کمکی است که تابع زیان را قبل از گام رو به عقب مقیاس‌بندی کرده و گرادیان‌ها را پیش از گام بهینه‌ساز به مقیاس اولیه بازمی‌گرداند. این دو برای کار در کنار هم طراحی شده‌اند؛ استفاده از یکی بدون دیگری خطای پیکربندی است و تست‌ها باید جلوی آن را بگیرند.

در این درس از autocast پردازنده (CPU) استفاده می‌کنیم تا در بستر CI قابل‌اجرا باشد؛ همین الگو با تغییر device_type="cpu" به device_type="cuda" عیناً روی GPU هم کار می‌کند. ابزار GradScaler روی CPU جنبه‌ی صوری (stub) دارد (زیرا autocast پردازنده به‌طور پیش‌فرض با BF16 کار می‌کند و نیازی به مقیاس‌بندی تابع زیان ندارد)، اما فراخوانی‌ها را داخل کد آورده‌ایم تا ساختار اتصالات دقیقاً با حلقه‌ی GPU یکسان باشد.

تشخیص NaN و Inf

این بررسی در دو نقطه انجام می‌شود. اول، خود تابع زیان پیش از گام رو به عقب با torch.isfinite بررسی می‌شود؛ زیانی که Inf یا NaN باشد گرادیان مفیدی تولید نمی‌کند و بدون ورود به بهینه‌ساز رد می‌شود. دوم، پس از اجرای scaler.unscale_(optimizer)، گرادیان‌های از-مقیاس‌افتاده با has_non_finite_grad(...) پیمایش می‌شوند و هر مقدار Inf یا NaN باعث رد شدن گام می‌گردد. ترکیب این دو بررسی، حالت‌های شکست (failure modes) در هر دو مرحله‌ی گذر رو به جلو و گام رو به عقب را پوشش می‌دهد.

خطایابی ضریب مقیاس‌بندی

ضریب مقیاس‌بندی همان وضعیت داخلی GradScaler است. در هر گام، مقدار scaler.get_scale() را خوانده و کنار نرخ یادگیری (learning rate) و نرم گرادیان لاگ می‌کنیم. در یک فرآیند آموزش سالم، ضریب مقیاس‌بندی با توان‌هایی از 2 بالا می‌رود تا جایی که حوالی 2^17 یا 2^18 اشباع شود. در فرآیندی که درست کار نمی‌کند، این ضریب مدام بین مقادیر بالا و پایین نوسان می‌کند؛ نشانه‌ای از این‌که گرادیان‌های مدل گاهی در بازه قرار می‌گیرند و گاهی نه. این ابزار تشخیصی بدون ثبت گزارش عملاً نامرئی می‌ماند.

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

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

• clip_global_l2_norm: پوششی دور torch.nn.utils.clip_grad_norm_ که نرم گرادیان را هم قبل از برش و هم بعد از آن برمی‌گرداند.
• has_non_finite_grad: تابع کمکی برای جست‌وجوی NaN و Inf در گرادیان‌ها.
• AmpTrainState: دربرگیرنده‌ی مدل، بهینه‌ساز AdamW، شیء GradScaler و دستگاه autocast. این کلاس متد step(inputs, targets) را در اختیار می‌گذارد که کل پایپ‌لاین (pipeline) برش، مقیاس‌بندی و رد کردن در صورت وجود NaN را اجرا می‌کند.
• StepLog و SkipLog: رکوردهای ساختاریافته برای لاگ هر گام.
• یک نمونه‌ی نمایشی که مدل کوچک nn.Linear را برای 20 گام آموزش می‌دهد، در گام 5 مقدار Inf به گرادیان تزریق می‌کند تا مسیر skip فعال شود و در نهایت لاگ خروجی را چاپ می‌کند.

اجرای کد:

python3 code/main.py

اسکریپت با کد خروج صفر به پایان می‌رسد و لاگ گام‌به‌گام را چاپ می‌کند که هر سطر آن با برچسب STEP یا SKIP مشخص شده است؛ حداقل یک سطر باید برچسب SKIP خورده باشد.

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

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

شمارنده‌ی گام‌های ردشده به عنوان هشدار، نه صرفاً یک سطر لاگ. چند گام ردشده در کل یک فرآیند آموزشی طبیعی است. اما صدها گام ردشده در هر دور (epoch) یک هشدار جدی است: مدل وارد وضعیتی شده که FP16 توان نگهداری مقادیر آن را ندارد و حلقه در سکوت با شکست مواجه می‌شود. در این درس نرخ رد شدن را در یک پنجره‌ی متحرک 1,000 گامی پایش می‌کنیم و در محیط عملیاتی، برای نرخ‌های بالای 5 درصد هشدار (page) ارسال می‌کنیم.

قرار دادن آستانه‌ی برش در فایل پیکربندی. مقدار max_norm = 1.0 مقدار پیش‌فرض امروزی برای آموزش مدل‌های زبانی است. ابتدا آن را روی یک مدل کوچک ارزیابی کنید؛ آستانه‌های بزرگ‌تر به مدل اجازه می‌دهند از بچ‌های واقعاً دشوار عبور کند؛ آستانه‌های کوچک‌تر بدترین حالت را مهار می‌کنند اما نوسان منحنی تابع زیان را بیشتر می‌کنند. این آستانه باید در همان فایل پیکربندی YAML یا JSON کنار زمان‌بند درس 44 قرار بگیرد.

ثبت نرم گرادیان در فایل CSV کنار مقادیر زمان‌بند. ستون‌های فایل CSV عبارتند از: step, lr, grad_l2_pre_clip, grad_l2_post_clip, loss, skipped, skip_reason, scaler_scale. هر ناظری با باز کردن این فایل، وضعیت زمان‌بندی، روند تغییرات گرادیان، ضریب مقیاس‌بندی و وضعیت اجرای گام (همراه با دلیل رد شدن احتمالی) را در یک سطر می‌بیند. پخش کردن این ستون‌ها در فایل‌های مختلف تحلیل‌ها را ناهماهنگ و سخت می‌کند.

اجرای scaler.update() در تمام گام‌ها، حتی در صورت رد شدن. در یک گام بی‌نقص، مقیاس‌بند شمارنده‌ی گام‌های بدون Inf خود را می‌خواند، یکی به آن اضافه می‌کند و در صورت نیاز ضریب را دو برابر می‌سازد. در گامی که رد شده باشد، مقیاس‌بند ضریب را نصف کرده و شمارنده را ریست می‌کند. فراموش کردن متد update() در مسیر گام ردشده همان باگی است که باعث می‌شود «ضریب مقیاس‌بندی هرگز تغییر نکند».

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

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

• تطابق دستگاه در autocast و بهینه‌ساز. برای آموزش روی GPU از torch.amp.autocast(device_type="cuda") و برای CPU از torch.amp.autocast(device_type="cpu") استفاده کنید. ناسازگاری دستگاه‌ها خطای نوع (type error) پنهانی ایجاد می‌کند؛ وضعیتی که در آن نمودار تابع زیان خوب به نظر می‌رسد اما مدل عملاً هیچ چیزی یاد نمی‌گیرد.
• بررسی تابع زیان پیش از گام رو به عقب. دستور torch.isfinite(loss).all() تنها یک عملیات تقلیل (reduction) روی تنسور است؛ هزینه‌ی محاسباتی آن ناچیز است اما با تشخیص بهینه‌ی زیان NaN، یک گام کامل آموزشی را نجات می‌دهد. همیشه آن را اجرا کنید.
• استفاده از set_to_none=True در zero_grad. این پارامتر گرادیان‌ها را به جای صفر روی None تنظیم می‌کند؛ در نتیجه بهینه‌ساز از محاسبات مربوط به گروه‌های پارامتری دست‌نخورده صرف‌نظر می‌کند. این تنظیم، هم نرخ پردازش را رایگان بالا می‌برد و هم احتمال بروز باگ را تا حدی کاهش می‌دهد.

خروجی این درس

در یک پروژه‌ی واقعی، فایل outputs/skill-clip-amp.md توضیح می‌دهد که این گام آموزش از چه آستانه‌ی برشی و چه دستگاه autocastای استفاده می‌کند، فایل CSV گام‌به‌گام در کجای سیستم کنترل نسخه (Version Control) نگهداری می‌شود و آستانه‌ی هشدار نرخ رد شدن در محیط عملیاتی چقدر است. این درس، موتور محرک این جریان را می‌سازد.

تمرین‌ها

1. تزریق دستی Inf را با یک جهش واقعی در تابع زیان (ضرب مقادیر هدف یک بچ در 1e8) جایگزین کنید و مطمئن شوید مسیر skip اجرا می‌شود.
2. حالتی با فلگ --bf16 اضافه کنید که به‌جای FP16، از BF16 در autocast استفاده کند. مقدار BF16 دامنه‌ی توان وسیع‌تری نسبت به FP16 دارد و به‌ندرت به مقیاس‌بندی تابع زیان نیازمند می‌شود؛ بررسی کنید که نرخ گام‌های ردشده در همین دمو به صفر برسد.
3. یک تست واحد (unit test) بنویسید که مطمئن شود تابع پوششی برش گرادیان، در حالتی که برشی رخ نمی‌دهد، نرم‌های قبل و بعد از برش را به‌درستی برمی‌گرداند.
4. قابلیت محاسبه‌ی نرخ رد شدن در یک پنجره‌ی متحرک را پیاده کرده و یک فلگ خط فرمان اضافه کنید که اگر این نرخ برای 100 گام متوالی از آستانه فراتر رفت، اجرای فرآیند را با خطا متوقف کند.
5. حلقه را طوری تنظیم کنید که فایل استاندارد CSV با ستون‌های (step, lr, grad_l2_pre_clip, grad_l2_post_clip, loss, skipped, skip_reason, scaler_scale) را ذخیره کند و مطمئن شوید با اعمال flush بعد از هر سطر، در صورت فشردن Ctrl-C اطلاعات ذخیره‌شده از بین نروند.

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

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

• Micikevicius et al., Mixed Precision Training (arXiv 1710.03740) - مقاله‌ی اصلی پیشنهاددهنده‌ی مقیاس‌بندی تابع زیان (Loss scaling)
• Pascanu, Mikolov, Bengio, On the difficulty of training recurrent neural networks (arXiv 1211.5063) - مقاله‌ی مرجع تکنیک برش گرادیان
• PyTorch torch.amp.GradScaler - مستندات API ابزار مقیاس‌بندی که در این درس پیاده شد
• PyTorch torch.nn.utils.clip_grad_norm_ - جزء سازنده‌ی (primitive) برش گرادیان مورد استفاده در این درس
• فاز 19 · 42 - ماژول دانلودی که پیکره‌ی متنی ورودی حلقه را تأمین می‌کند
• فاز 19 · 43 - دیتالودری که این حلقه از آن استفاده می‌کند
• فاز 19 · 44 - زمان‌بندی که این حلقه با آن ترکیب می‌شود