چک‌پوینت‌گذاری گرادیان و بازسازی اکتیویشن در ترنسفورمرها

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

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

چک‌پوینت‌گذاری گرادیان (Gradient Checkpointing یا Activation Recomputation) این مسئله را با معاوضه‌ی کنترل‌شده‌ی محاسبات و حافظه حل می‌کند: به‌جای نگه‌داشتن همه‌ی اکتیویشن‌ها در حافظه، بخش زیادی از آن‌ها دور ریخته شده و در گام بازگشت دوباره محاسبه می‌شوند. در این درس یاد می‌گیرید چطور با استراتژی‌های گزینشی، حافظه را تا 5 برابر آزاد کنید بدون اینکه سربار زمانی سنگینی بپردازید.

در روش پس‌انتشار، تمام اکتیویشن‌های میانی ذخیره می‌شوند. در مدلی با 70 میلیارد پارامتر و طول بافت 128K، این مقدار به 3 ترابایت اکتیویشن به ازای هر رتبه (Rank) می‌رسد. چک‌پوینت‌گذاری با فدا کردن محاسبات (FLOPs)، حافظه آزاد می‌کند: محاسبه‌ی دوباره به‌جای ذخیره‌سازی. مسئله‌ی اصلی این است که کدام بخش‌ها را حذف کنیم، و پاسخ قطعاً «همه‌ی آن‌ها» نیست.

صورت مسئله

هنگام آموزش ترنسفورمر (Transformer)، ورودیِ تمام عملیات‌هایی که در گام رو به عقب (Backward Pass) مشتق‌گیری می‌شوند برای هر لایه ذخیره خواهد شد: ورودی‌های اتنشن (Attention)، تصویر کردن‌های Q/K/V، خروجی سافت‌مکس (Softmax)، ورودی‌های FFN، خروجی‌های نرمال‌سازی و مسیر باقیمانده (Residual stream). برای لایه‌ای با اندازه‌ی مخفی d، طول توالی L و اندازه‌ی دسته (Batch) برابر با B، این حجم تقریباً معادل 12 * B * L * d عدد اعشاری به ازای هر لایه است.

برای d=8192, L=8192, B=1، این رقم به حدود 805 میلیون مقدار می‌رسد که در دقت BF16 معادل 1.6 گیگابایت به ازای هر لایه است. یک مدل 64 لایه‌ای تقریباً 103 گیگابایت اکتیویشن خواهد داشت؛ آن هم پیش از ضرب در اندازه‌ی میکرودسته، پیش از اضافه شدن مقادیر میانی اتنشن-سافت‌مکس (L^2 به ازای هر سر)، و پیش از به حساب آوردن کپی‌های جزئی ناشی از موازی‌سازی تنسور (Tensor Parallelism).

این صورت‌حساب دوطرفه است: وزن‌های BF16 به‌همراه وضعیت بهینه‌ساز شاید در 80 گیگابایت جا شوند، اما حجم اکتیویشن‌ها حافظه را پر می‌کند و شما را فراتر از آن می‌برد. راهکار استاندارد این مسئله، چک‌پوینت‌گذاری گرادیان (Gradient Checkpointing که بازسازی اکتیویشن یا Activation Recomputation نیز نامیده می‌شود) است. بیشتر اکتیویشن‌ها را دور بریزید؛ در گام رو به عقب، گذر رو به جلو (Forward Pass) را دوباره اجرا کنید تا بازسازی شوند. هزینه‌ی کار: FLOPs اضافی. مزیت آن: کاهش حافظه به نسبتِ تعداد قطعات چک‌پوینت به کل لایه‌ها.

اگر چک‌پوینت‌گذاری ساده‌لوحانه پیاده شود، در هر گام حدود 33% محاسبات گذر رو به جلوی بیشتری می‌طلبد. اما اگر درست و هوشمندانه پیاده شود — یعنی چک‌پوینت‌گذاری گزینشی مطابق «انتخاب هوشمندِ» Korthikanti و همکاران — حافظه 5 برابر سبک‌تر می‌شود، در حالی که سربار محاسباتی (FLOP overhead) زیر 5% است. این موضوع با وجود ضرب‌های ماتریسی FP8، تخلیه روی پردازنده‌ی اصلی (CPU Offload) در FSDP و معماری MoE با موازی‌سازی خبره (Expert Parallelism) اهمیت دوچندانی پیدا می‌کند؛ چون در چنین شرایطی نه حافظه‌ی اضافه‌ای دارید و نه توان پردازشی هرز.

مفاهیم پایه

گام رو به عقب واقعاً به چه چیزهایی نیاز دارد؟

فرض کنید output = layer(input). گام رو به عقب به دنبال grad_input و grad_params است. برای محاسبه‌ی آن‌ها به این موارد نیاز دارد:

• خود input (برای محاسبه‌ی grad_params = input.T @ grad_output در لایه‌های خطی)؛
• برخی از مقادیر میانی مشتق اکتیویشن‌ها (مشتق ReLU/GELU/سافت‌مکس به مقدار خود اکتیویشن وابسته است).

گذر رو به جلو این مقادیر را خودکار در گراف مشتق‌گیری خودکار (Autograd) ذخیره می‌کند. هر فراخوانی tensor.retain_grad() و هر عملیاتی که به ورودی‌اش نیاز داشته باشد، ارجاعی به آن نگه می‌دارد.

چک‌پوینت‌گذاری کامل و خام (Naive Full Checkpointing)

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

مثال: ترنسفورمر 32 لایه‌ای که به 32 قطعه‌ی 1 لایه‌ای تقسیم شده است.

• مصرف حافظه: ورودیِ 32 لایه (خیلی کم) در برابر 32 ضرب‌در حجم اکتیویشن هر لایه (بسیار زیاد).
• بار محاسباتی اضافه: 1 گذر رو به جلوی اضافی به ازای هر قطعه؛ یعنی حدود 33% افزایش در کل FLOPs گذر رو به جلو (از آنجا که گام رو به عقب معمولاً 2 برابرِ گذر رو به جلو محاسبات دارد، کل گام آموزشی به‌جای 1 + 2 = 3 واحد، برابر با 1 + 1 + 2 = 4 واحد محاسباتی می‌شود).

این همان فرمول اصلیِ مقاله‌ی Chen و همکاران در سال 2016 است: 1 چک‌پوینت به ازای هر sqrt(L) لایه برای ایجاد تعادل میان حافظه و محاسبات. برای $L=64$، یعنی 8 چک‌پوینت.

چک‌پوینت‌گذاری گزینشی (Selective Checkpointing - Korthikanti 2022)

هزینه‌ی همه‌ی اکتیویشن‌ها یکسان نیست. خروجی سافت‌مکس اتنشن برابر با B*L*L*heads است و با افزایش طول توالی به‌صورت درجه دو رشد می‌کند. در مقابل، اکتیویشن پنهان FFN برابر با B*L*4d است و رشد خطی دارد. در توالی‌های طولانی، سافت‌مکس عملاً حافظه را می‌بلعد.

در چک‌پوینت‌گذاری گزینشی، اکتیویشن‌های کم‌هزینه از نظر حافظه (تصویر کردن‌های خطی، مسیرهای باقیمانده) حفظ می‌شوند و تنها اکتیویشن‌های پرهزینه (اتنشن) دوباره محاسبه خواهند شد. این‌گونه با کمترین هزینه‌ی محاسباتی برای بازسازی، حافظه‌ی چشمگیری از $O(L^2)$ آزاد می‌شود.

مجموعه‌ی Megatron-Core این قابلیت را تحت عنوان بازسازی «گزینشی» اکتیویشن پیاده کرده است که در اکثر تمرین‌های مدل‌های پیشرو از سال 2024 به بعد به کار می‌رود.

تخلیه روی پردازنده‌ی اصلی (Offload)

رویکرد جایگزین بازسازی: اکتیویشن‌ها را در فاصله‌ی بین گذر رو به جلو و عقب، به رمِ پردازنده‌ی اصلی بفرستید. این روش به پهنای باند PCIe متکی است؛ هر زمان پهنای باند آزاد سیستم اقتصادی‌تر از هزینه‌ی بازسازی محاسباتی باشد، این رویکرد می‌درخشد. استراتژی‌های ترکیبی هم رایج‌اند: برای برخی لایه‌ها چک‌پوینت بگذارید و بقیه را روی پردازنده‌ی اصلی تخلیه کنید.

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

مدل هزینه‌ی بازسازی

میزان FLOPs در هر گام با چک‌پوینت‌گذاری ساده در هر k لایه از مجموع L لایه:

flops_fwd_normal = L * f_layer
flops_bwd_normal = 2 * L * f_layer
flops_total_normal = 3 * L * f_layer

flops_fwd_ckpt = L * f_layer
flops_recompute = L * f_layer  # one extra forward per layer in the segment
flops_bwd_ckpt = 2 * L * f_layer
flops_total_ckpt = 4 * L * f_layer
overhead = 4 / 3 - 1 = 0.33 = 33%

در چک‌پوینت‌گذاری گزینشی، به‌جای کل لایه فقط کرنل (Kernel) اتنشن را دوباره محاسبه می‌کنید:

flops_recompute_selective = L * f_attention ~= L * f_layer * 0.15
overhead_selective = (3 + 0.15) / 3 - 1 = 0.05 = 5%

مدل صرفه‌جویی در حافظه

حجم اکتیویشن در هر لایه: A. برای L لایه، کل حافظه‌ی اکتیویشن برابر است با: L * A.

چک‌پوینت‌گذاری کامل (اندازه‌ی قطعه برابر با 1): فقط L * input_volume ذخیره می‌شود (حدود L * 1/10 A در یک ترنسفورمر استاندارد). صرفه‌جویی تقریباً معادل 9 * L * A * 1/10.

چک‌پوینت‌گذاری در هر k لایه: ذخیره‌ی L/k * A به‌علاوه‌ی حافظه‌ی مربوط به k-1 لایه درون قطعه‌ی فعال.

در k = sqrt(L)، هزینه‌ی حافظه و هزینه‌ی بازسازی هر دو متناسب با sqrt(L) مقیاس می‌شوند — این بهینه‌ترین نقطه‌ی تعادل برای لایه‌هایی با هزینه‌ی یکسان است.

چه زمانی نباید چک‌پوینت بگذاریم؟

• لایه‌های درونی در یکی از مراحل پایپ‌لاین (Pipeline) که پردازش آن‌ها در جریان است؛ چون به‌هر‌حال باید کارشان تمام شود.
• لایه‌های اول و آخر، البته اگر بیشترین بار محاسباتی مرحله را به دوش بکشند (که در ترنسفورمرها پدیده‌ای نادر است).
• کرنل‌های اتنشنی که از FlashAttention استفاده می‌کنند؛ چون این الگوریتم خودش سافت‌مکس را سریع بازسازی می‌کند و در نتیجه چک‌پوینت‌گذاریِ بیشتر در سطح لایه فایده‌ی محسوسی ندارد.

الگوهای پیاده‌سازی

1. پوشش‌دهنده‌ی تابع (Function wrapper): یک قطعه را درون torch.utils.checkpoint.checkpoint(fn, input) بپیچید. PyTorch فقط input را نگه می‌دارد و سایر بخش‌ها را در گام رو به عقب دوباره محاسبه می‌کند.

2. مبتنی بر دکوراتور (Decorator-based): لایه‌ها را به‌عنوان قابل‌چک‌پوینت علامت‌گذاری کنید؛ سیستم ترین در زمان اعمال پیکربندی تصمیم می‌گیرد کدام قطعات پوشش داده شوند.

3. محاسبه‌ی دوباره‌ی دستی و صریح: گام رو به عقب را خودتان بنویسید و تابع سفارشی recompute_forward را صدا بزنید تا گذر رو به جلو را با همان ورودی ذخیره‌شده بازآفرینی کند.

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

تداخل با TP ،PP و FP8

• موازی‌سازی تنسور (Tensor Parallel): ورودی‌های چک‌پوینت باید در زمان بازسازی جمع‌آوری یا دوباره پخش شوند؛ حواستان به هزینه‌ی ارتباطات شبکه باشد.
• موازی‌سازی خط لوله‌ای (Pipeline Parallel): الگوی رایج این است که گذر رو به جلوی هر مرحله از پایپ‌لاین چک‌پوینت شود تا میکرودسته‌ها در حرکت معکوس بتوانند از حافظه‌ی اکتیویشن دوباره استفاده کنند.
• بازسازی در FP8: تاریخچه‌ی amax که حین بازسازی به‌روزرسانی می‌شود باید دقیقاً با گذر رو به جلوی اولیه مطابقت داشته باشد؛ در غیر این صورت مقیاس FP8 دچار انحراف می‌شود. بیشتر فریم‌ورک‌ها از این مقیاس یک اسنپ‌شات (Snapshot) می‌گیرند.

پیاده‌سازی گام‌به‌گام

مرحله‌ی 1: یک مدل اسباب‌بازی به‌همراه قطعات

import numpy as np


def linear_forward(x, w, b):
    return x @ w + b


def relu(x):
    return np.maximum(x, 0)


def layer_forward(x, w1, b1, w2, b2):
    h = relu(linear_forward(x, w1, b1))
    return linear_forward(h, w2, b2)


def model_forward(x, params):
    activations = [x]
    h = x
    for w1, b1, w2, b2 in params:
        h = layer_forward(h, w1, b1, w2, b2)
        activations.append(h)
    return h, activations

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

def model_backward(grad_output, activations, params):
    grads = [None] * len(params)
    g = grad_output
    for i in range(len(params) - 1, -1, -1):
        w1, b1, w2, b2 = params[i]
        x_in = activations[i]
        h_pre = linear_forward(x_in, w1, b1)
        h = relu(h_pre)
        gh = g @ w2.T
        gw2 = h.T @ g
        gb2 = g.sum(axis=0)
        g_pre = gh * (h_pre > 0)
        gx = g_pre @ w1.T
        gw1 = x_in.T @ g_pre
        gb1 = g_pre.sum(axis=0)
        grads[i] = (gw1, gb1, gw2, gb2)
        g = gx
    return g, grads

مرحله‌ی 3: حافظه‌ی لازم در استراتژی چک‌پوینت در هر k لایه

def model_forward_checkpointed(x, params, k=4):
    saved_inputs = [x]
    h = x
    for i, (w1, b1, w2, b2) in enumerate(params):
        h = layer_forward(h, w1, b1, w2, b2)
        if (i + 1) % k == 0:
            saved_inputs.append(h)
    return h, saved_inputs


def model_backward_checkpointed(grad_output, saved_inputs, params, k=4):
    grads = [None] * len(params)
    g = grad_output
    segments = [(j * k, min((j + 1) * k, len(params))) for j in range(len(saved_inputs))]
    for seg_idx in range(len(saved_inputs) - 1, -1, -1):
        start, end = segments[seg_idx]
        if start >= end:
            continue
        x_in = saved_inputs[seg_idx]
        _, seg_acts = model_forward(x_in, params[start:end])
        g, seg_grads = model_backward(g, seg_acts, params[start:end])
        for j, gr in enumerate(seg_grads):
            grads[start + j] = gr
    return g, grads

مرحله‌ی 4: مدل هزینه‌ی محاسباتی

def checkpoint_cost(n_layers, segment_size, flops_per_layer=1.0):
    fwd = n_layers * flops_per_layer
    recompute = n_layers * flops_per_layer
    bwd = 2 * n_layers * flops_per_layer
    return {
        "fwd": fwd,
        "recompute": recompute,
        "bwd": bwd,
        "total": fwd + recompute + bwd,
        "overhead_vs_no_ckpt": (fwd + recompute + bwd) / (fwd + bwd) - 1.0,
    }


def selective_checkpoint_cost(n_layers, attention_fraction=0.15,
                              flops_per_layer=1.0):
    fwd = n_layers * flops_per_layer
    recompute = n_layers * attention_fraction * flops_per_layer
    bwd = 2 * n_layers * flops_per_layer
    return {
        "fwd": fwd,
        "recompute": recompute,
        "bwd": bwd,
        "total": fwd + recompute + bwd,
        "overhead_vs_no_ckpt": (fwd + recompute + bwd) / (fwd + bwd) - 1.0,
    }

مرحله‌ی 5: تخمین‌گر حافظه

def activation_memory_mb(n_layers, hidden=8192, seq=8192,
                        batch=1, bytes_per_value=2):
    per_layer = 12 * batch * seq * hidden * bytes_per_value
    return n_layers * per_layer / 1e6


def memory_after_checkpoint(n_layers, segment_size, hidden=8192,
                           seq=8192, batch=1, bytes_per_value=2):
    n_seg = max(1, n_layers // segment_size)
    saved = (n_seg + segment_size) * 1 * batch * seq * hidden * bytes_per_value
    return saved / 1e6

مرحله‌ی 6: اندازه‌ی بهینه‌ی قطعه

def optimal_segment(n_layers):
    return int(round(np.sqrt(n_layers)))

مرحله‌ی 7: تصمیم‌گیری درباره‌ی چک‌پوینت‌گذاری گزینشی

def should_recompute(layer_type, activation_bytes, recompute_flops_ratio):
    if layer_type == "attention" and activation_bytes > 100 * 1e6:
        return True
    if layer_type == "ffn" and activation_bytes > 500 * 1e6:
        return recompute_flops_ratio < 0.1
    return False

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

• torch.utils.checkpoint: با دستور from torch.utils.checkpoint import checkpoint — پوشش‌دهنده‌ی استاندارد در PyTorch. یک تابع را ورودی می‌گیرد، فقط ورودی‌ها را ذخیره می‌کند و در گام رو به عقب دوباره محاسبات را پیش می‌برد.
• بازسازی اکتیویشن در Megatron-Core: از حالت‌های selective، full و block پشتیبانی می‌کند. رویکرد مرسوم در آموزش مدل‌های مرزی از سال 2024 به این سو.
• تخلیه روی پردازنده‌ی اصلی در FSDP2: با متد module.to_empty(device="cpu") و تنظیم offload_policy در FSDP2، اکتیویشن‌ها به‌جای محاسبه‌ی دوباره روی پردازنده‌ی اصلی ذخیره می‌شوند.
• DeepSpeed ZeRO-Offload: تخلیه روی پردازنده‌ی اصلی برای وضعیت بهینه‌ساز و اکتیویشن‌ها؛ مکملی قدرتمند برای چک‌پوینت‌گذاری.

خروجی نهایی

در پایان این درس فایل outputs/prompt-activation-recompute-policy.md تولید می‌شود — پرامپتی (Prompt) که مشخصات مدل شما (لایه‌ها، بعد پنهان، طول توالی و اندازه دسته) و حافظه‌ی آزاد GPU را می‌گیرد و یک سیاست بازسازی لایه‌به‌لایه تحویلتان می‌دهد (هیچ / گزینشی / کامل / تخلیه روی پردازنده‌ی اصلی).

تمرین‌ها

1. درستی خروجی را بسنجید. ترکیب model_forward + model_backward (با اکتیویشن کامل) را با model_forward_checkpointed + model_backward_checkpointed (با قطعات) مقایسه کنید. گرادیان‌های پارامترها باید تا حد دقت ماشین کاملاً یکسان باشند.

2. اندازه‌ی قطعه‌ی k را از 1 تا L پیمایش کنید. بار محاسباتی اضافه و حافظه را رسم کنید و نقطه‌ی زانویی (Knee) نمودار را به دست آورید.

3. چک‌پوینت‌گذاری گزینشی را پیاده کنید: ورودی ماژول اتنشن را نگه دارید اما مقادیر میانی درون آن را دور بریزید. بار محاسباتی اضافه را در مقایسه با چک‌پوینت کل لایه برای مدلی 32 لایه‌ای در طول توالی 8192 بسنجید.

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

5. عملکرد یک ترنسفورمر واقعی را در PyTorch با و بدون torch.utils.checkpoint بسنجید. حافظه‌ی مصرفی را با torch.cuda.max_memory_allocated و زمان سپری‌شده برای هر گام را ثبت کنید.

واژگان کلیدی

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

• Chen et al., 2016 -- "Training Deep Nets with Sublinear Memory Cost" — مقاله‌ی اصلی که چک‌پوینت‌گذاری گرادیان را به شکل رسمی تبیین کرد.
• Korthikanti et al., 2022 -- "Reducing Activation Recomputation in Large Transformer Models" — بازسازی گزینشی اکتیویشن و تحلیل دقیق هزینه‌های آن.
• Pudipeddi et al., 2020 -- "Training Large Neural Networks with Constant Memory using a New Execution Algorithm" — رویکرد جایگزین حافظه-ثابت از طریق مادیت‌بخشی در حالت معکوس.
• Ren et al., 2021 -- "ZeRO-Offload: Democratizing Billion-Scale Model Training" — تخلیه اکتیویشن‌ها روی پردازنده‌ی اصلی در مقیاس‌های کلان.
• PyTorch torch.utils.checkpoint docs — مستندات رابط برنامه‌نویسی استاندارد در PyTorch.
• Megatron Bridge activation recomputation documentation — مستندات حالت‌های selective ،full و block.