چیت‌شیت
حافظه آموزش، انباشت و چک‌پوینت‌ها
Training Memory, Accumulation and Checkpoints
رفع گام‌های NaN، محدودیت‌های حافظه، مقیاس انباشت و ادامه‌های ناموفق. کلیپ، AMP، باز‌محاسبه، چک‌پوینت‌های اتمیک.
فهرست مطالب
نشانه‌های آموزش
عیبیابیصعود عمودی تابع زیان (loss) پس از یک دسته
- علت
- جهش نُرم گرادیان، بدون کلیپ
- راه‌حل
clip_grad_norm_درmax_norm = 1.0
NaN شدن تمام وزن‌ها پس از یک گام
- علت
- سرریز FP16 مقدار Inf و سپس NaN می‌دهد
- راه‌حل
torch.amp.GradScalerگام را نادیده می‌گیرد
فاکتور مقیاس‌بندی هرگز تغییر نمی‌کند
- علت
- عدم وجود
scaler.update()در مسیر رد شدن - راه‌حل
- فراخوانی
scaler.update()در هر گام
اعمال کلیپ روی گرادیان‌های مقیاس‌شده
- علت
- اجرای کلیپ پیش از unscale
- راه‌حل
- ابتدا
scaler.unscale_(optimizer)را صدا بزنید
صدها رد شدن (skip) در هر اپوک
- علت
- مدل فراتر از بازه FP16 است
- راه‌حل
- هشدار در صورت نرخ رد شدن بالای 5 درصد
گام بهینه‌ساز 16 برابر بزرگ‌تر است
- علت
- تابع زیان (loss) میکروبچ تقسیم نشده است
- راه‌حل
- تقسیم تابع زیان بر
accum_steps
تفاوت در منحنی تابع زیان (loss) پس از ادامه
- علت
- نبود وضعیت RNG در چک‌پوینت
- راه‌حل
- ذخیره خروجی
capture_rng_state()
اکتیویشن‌ها در حافظه جا نمی‌شوند
- علت
- پس‌انتشار همه اکتیویشن‌ها را نگه می‌دارد
- راه‌حل
torch.utils.checkpoint.checkpoint(fn, input)
خطاهای بارگذاری چک‌پوینت
عیب‌یابیworld_size mismatch- علت
- ادامه روی یک اندازه‌ی جهان متفاوت
- راه‌حل
- ادامه روی همان اندازه‌ی جهان
sha256 mismatch on rank- علت
- ناقص شدن شارد حین flush
- راه‌حل
- نوشتن در tmp.، سپس fsync و تغییر نام
shard file missing- علت
- فایل‌های کمتر از رتبه‌ها نسبت به شاردهای مانیفست
- راه‌حل
- بررسی و پیمایش تک‌تک شاردها
unknown schema- علت
- اسکیمای پی‌لود با ckpt آغاز نمی‌شود
- راه‌حل
- انتقال‌های branch بر اساس رشته اسکیما
NaN شدن تابع زیان (loss) پس از 100 گام از ادامه
- علت
- شارد ناقص در پشت مانیفست موجود
- راه‌حل
- نوشتن اتمیک به همراه sha256 حین بارگذاری
پر شدن دیسک در میانه اجرا
- علت
- عدم وجود گردش چک‌پوینت‌ها
- راه‌حل
rotate_checkpointsآخرین K مورد را نگه می‌دارد
گام امن AMP
کدpython
optimizer.zero_grad(set_to_none=True) # تبدیل گرادیان‌ها به None
with torch.amp.autocast(device_type="cuda"): # گذر رفت FP16
loss = loss_fn(model(inputs), targets)
if not torch.isfinite(loss).all().item(): # بررسی پیش از پس‌انتشار
return skip("non_finite_loss") # بدون scaler.update در اینجا
scaler.scale(loss).backward() # پس‌انتشار مقیاس‌شده
scaler.unscale_(optimizer) # unscale پیش از کلیپ
if has_non_finite_grad(model.parameters()):
scaler.update() # فاکتور را نصف می‌کند
return skip("non_finite_grad")
pre, post = clip_global_l2_norm(params, max_norm)
scaler.step(optimizer)
scaler.update() # دو برابر یا نصف می‌کند
انباشت، همگام‌سازی در آخرین مرحله
کدpython
accum_steps = len(micro_batches)
zero_grads(model)
for i, (x, y) in enumerate(micro_batches):
is_last = i == accum_steps - 1
if no_sync_until_last and not is_last:
with no_sync_context(model): # DDP عملیات all-reduce را رد می‌کند
loss = loss_scaled_for_accum(model(x), y, accum_steps, loss_fn)
loss.backward() # گرادیان‌ها به param.grad افزوده می‌شوند
else:
loss = loss_scaled_for_accum(model(x), y, accum_steps, loss_fn)
loss.backward() # یک همگام‌سازی به ازای هر گام مؤثر
optimizer.step() # یک‌بار به ازای هر دسته‌ی مؤثر
وظایف در کد
کار و کد| وظیفه | کد |
|---|---|
| کلیپ بر اساس نرم L2 سراسری | clip_grad_norm_(parameters, max_norm) |
| ثبت لاگ فاکتور مقیاس‌بندی | scaler.get_scale() |
| مقیاس‌بندی تابع زیان (loss) میکروبچ | loss_scaled_for_accum(logits, target, accum_steps, loss_fn) |
| باز‌محاسبه یک قطعه در پس‌انتشار | checkpoint(fn, input) |
| ذخیره کامل وضعیت به صورت اتمیک | save_checkpoint(model, optimizer, scheduler, state, out_path) |
| بارگذاری یک فایل غیرقابل اعتماد | torch.load(path, map_location="cpu", weights_only=True) |
| ذخیره و بارگذاری شاردهای به ازای هر رتبه | save_sharded(per_rank_state, dest_dir, step) |
| ادامه آموزش و راستی‌آزمایی هر شارد | load_sharded(src_dir, expected_world_size) |
ذخیره‌ی اتمیک شاردبندی‌شده
معماریذخیره و ادامه از میانه اپوک
گام‌به‌گام- save_checkpointبسته‌بندی مدل، بهینه‌ساز، اسکژولر، شمارنده‌ها و RNG.
- atomic_saveیک فایل موقت می‌نویسد، سپس
os.replaceرا صدا می‌زند. - load_checkpointاجرای
torch.loadباweights_only=True. - restore_rng_stateبازگردانی RNG در پایتون، NumPy و torch.
- حلقه آموزشادامه از
batch_in_epoch. تفاوت‌های بازنگری زیان زیر 1e-4.
اجرای درس‌ها
دستورها| دستور | کارکرد |
|---|---|
python3 code/main.py | درس 45: 20 گام، ردیف‌های برچسب‌خورده با STEP یا SKIP |
--num-steps 100 | درس 46: پیمایش طولانی‌تر برای رسم منحنی |
--accum-grid 1,2,4,8,16 | درس 46: گام‌های انباشت جهت پیمایش |
--total-steps 24 --interrupt-at 10 | درس 47: پنجره دمو برای ادامه کار |
python3 code/main.py | درس 80: 4 شارد به علاوه مانیفست، بارگذاری مجدد بایت‌به‌بایت برابر |
python3 code/main.py | درس 81: 20 گام، 4 رتبه، RESUME VERIFIED |
پیش‌فرض‌ها در کد
پیش‌فرض‌ها| تنظیم | پیش‌فرض | نکته |
|---|---|---|
DEFAULT_MAX_NORM | 1.0 | آستانه کلیپ در AmpTrainState |
`amp_dtype` on cpu | torch.bfloat16 | torch.float16 روی cuda |
window | 1000 | گام‌ها در rolling_skip_rate |
k | 4 | لایه‌ها به ازای هر قطعه چک‌پوینت |
attention_fraction | 0.15 | سهم باز‌محاسبه، حالت گزینشی |
keep_last | 5 | چک‌پوینت‌هایی که rotate_checkpoints نگه می‌دارد |
CHECKPOINT_SCHEMA | ckpt.v1 | رشته اسکیما در پی‌لود |
CHECKPOINT_STEP | 10 | گام ذخیره شاردبندی‌شده درس 81 |
اعداد مربوط به حافظه و هزینه
مفهوم‌ها- اکتیویشن‌ها به ازای هر لایه
- حدود
12 * B * L * dعدد اعشاری. - d=8192, L=8192, B=1
- حدود 1.6 گیگابایت به ازای هر لایه در BF16. شصت و چهار لایه حدود 103 گیگابایت.
- چک‌پوینت‌گیری کامل
- حدود 33٪ فلاپس بیشتر. یک گام به‌جای 3 واحد، 4 واحد است.
- چک‌پوینت‌گذاری گزینشی
- فقط باز‌محاسبه اتنشن. حدود 5٪ سربار.
- قاعده sqrt(L)
- یک چک‌پوینت به ازای هر sqrt(L) لایه. 8 چک‌پوینت برای L=64.
- دسته‌ی مؤثر
micro_batch * accum_stepsضربدر اندازه‌ی جهان موازی‌سازی داده.- حافظه بهینه‌ساز ZeRO-1
- معادل 12P/N بایت به ازای هر رتبه.
پرسش و پاسخ
پرسش و پاسخ- چرا GradScaler پردازنده کاری انجام نمی‌دهد؟
- تبدیل خودکار CPU از BF16 استفاده می‌کند؛ نیازی به مقیاس‌بندی تابع زیان ندارد.
- فاکتور چه زمانی تغییر می‌کند؟
- پس از توالی‌های موفق دو برابر می‌شود. با هر رد شدن (skip) نصف می‌شود.
- آیا انباشت زمان را ذخیره می‌کند؟
- خیر. دو برابر کردن
accum_stepsزمان واقعی هر گام را دو برابر می‌کند. - چرا
weights_only=True؟ - فایل
.ptیک پیکل است. پیکل‌گشایی کد اجرا می‌کند. از torch نسخه 2.6 یا جدیدتر استفاده کنید. - چرا یک فایل به ازای هر رتبه؟
- پهنای باند نوشتن مقیاس‌پذیر می‌شود. 4 ساعت روی 1 رتبه، روی 64 رتبه 4 دقیقه است.
نکات
نکته‌ها- دستگاه autocast را با دستگاه بهینه‌ساز مطابقت دهید.
- به جای assert از
ValueErrorاستفاده کنید؛ فلگpython -Oگزاره‌های assert را حذف می‌کند. - در صورت جهش تابع زیان (loss) به بیش از 2 برابر، به چک‌پوینت قبلی بازگردید.