梯度检查点与激活重算

反向传播保留每个中间激活。在 70B 参数、128K context 下，那是每 rank 3 TB 的激活。检查点用 FLOPs 换内存：重算而不保存。问题是丢哪些段，答案不是 "全部"。

类型： Build 语言： Python（配合 numpy，可选 torch） 前置要求： 阶段 10 第 04 课（预训练 Mini-GPT），阶段 10 第 05 课（规模化与分布式） 预计时间： ~70 分钟

问题背景

训练一个 transformer 为每一层存反向时被微分的每个操作的输入：注意力输入、Q/K/V 投影、softmax 输出、FFN 输入、norm 输出和残差流。对一个 hidden size 为 d、序列长度 L、batch B 的层，这是每层约 12 * B * L * d 个 float 的量级。

对 d=8192, L=8192, B=1，那是 BF16 下每层 800 MB。一个 64 层模型是 51 GB 激活——而且这还没乘以微批次大小、没加上注意力 softmax 的中间量（每头 L^2）、没算上张量并行的部分拷贝。

两面夹击的账单：BF16 权重加优化器状态也许能塞进 80GB，但激活把你推过去了。梯度检查点（即激活重算）是标准修法。丢掉大部分激活；反向时重做前向把它们取回来。代价：额外 FLOPs。收益：内存按检查点段数对总层数的比例下降。

朴素地做，检查点每步大约多花 33% 的前向 FLOPs。做得好——按 Korthikanti et al. 的 "智能选择" 做选择性检查点——你以低于 5% 的 FLOP 开销省下 5 倍内存。而有了 FP8 matmul、FSDP 卸载和专家并行 MoE，这真的要紧：内存和浪费的算力你都付不起。

核心概念

反向实际需要什么

output = layer(input)。反向想要 grad_input 和 grad_params。要算它们它需要：

• input（对线性层算 grad_params = input.T @ grad_output）
• 一些激活导数中间量（ReLU/GELU/softmax 的导数取决于激活值）

前向传播自动把这些存进 autograd 图。每个 tensor.retain_grad() 和每个需要其输入的操作都保留一个引用。

朴素的全检查点

把网络切成 N 段。前向时，只存每段的 输入。反向需要中间量时，重跑那段的前向把它们物化出来，再微分。

例子：32 层 transformer 切成 32 段，每段 1 层。

• 内存：32 个层输入（小）vs 32 * （每层激活体积）（巨大）。
• 额外计算：每段多一次前向，即总共多约 33% 的前向 FLOPs（因为反向是 2 倍前向，完整步骤变成 1 + 1 + 2 = 4 个单位而不是 1 + 2 = 3）。

这是最初 Chen et al. 2016 的配方：每 sqrt(L) 层一个检查点来平衡内存和计算。对 L=64，那是 8 个检查点。

选择性检查点（Korthikanti 2022）

不是所有激活成本一样。注意力 softmax 输出是 B*L*L*heads，随序列长度 二次 增长。FFN 隐藏激活是 B*L*4d，线性增长。对长序列，softmax 主导。

选择性检查点保留存储便宜的激活（线性投影、残差），只重算昂贵的（注意力）。你付最少的 FLOPs 去重算，但省下 O(L^2) 的内存。

Megatron-Core 把这实现为 "选择性" 激活重算。在大多数 2024+ 前沿训练运行里使用。

卸载

重算的替代方案：在前向和反向之间把激活运到 CPU 内存。需要 PCIe 带宽；当空闲带宽超过重物化成本时有益。混合策略很常见：一些层检查点，另一些卸载。

FSDP2 把卸载作为一流选项。卸载在 GPU 被内存瓶颈卡住但 CPU-GPU 传输有余量时大放异彩。

重算成本模型

在 L 层中每 k 层做朴素检查点的每步 FLOPs：

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  # 段里每层一次额外前向
flops_bwd_ckpt = 2 * L * f_layer
flops_total_ckpt = 4 * L * f_layer
overhead = 4 / 3 - 1 = 0.33 = 33%

用选择性检查点你只重算注意力核，不是整层：

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（标准 transformer 约 L * 1/10 A）。省下约 9 * L * A * 1/10。

每 k 层检查点：存 L/k * A 加活跃段内 k-1 层的量。

在 k = sqrt(L) 时，内存和重算成本都随 sqrt(L) 缩放——对等成本层的最优权衡。

什么时候不检查点

• 一个已在途的流水线阶段的最内层。它们反正得完成。
• 如果第一层和最后一层主导该阶段的计算（在 transformer 里罕见）。
• 已经用 FlashAttention 的注意力核——Flash 已经快速重算 softmax 了，所以额外的层级检查点在它之上加不了多少。

实现模式

1. 函数封装： 把一段包进 torch.utils.checkpoint.checkpoint(fn, input)。PyTorch 只存 input，反向时重算其他一切。

2. 基于装饰器： 把层标为可检查点；训练器在配置时决定哪些段被包装。

3. 手动显式重算： 自己写反向传播，调一个用存下的输入复制前向的自定义 recompute_forward。

三者给出相同的功能结果。封装是标准惯用法。

和 TP / PP / FP8 的交互

• 张量并行： 检查点输入在重算时必须被收集或重新分散；处理通信成本。
• 流水线并行： 典型模式是检查点每个流水线阶段的前向，这样逆序微批次能复用激活内存。
• FP8 重算： 重算时更新的 amax 历史必须和原始前向的匹配，否则 FP8 缩放漂移。大多数框架快照这个缩放。

activation-recompute

动手构建

第 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 卸载：在 FSDP2 里用 module.to_empty(device="cpu") 配 offload_policy，把激活分片到 CPU 而不重算。
• DeepSpeed ZeRO-Offload：优化器状态和激活的 CPU 卸载，和检查点互补。

拿去用

本节课产出 outputs/prompt-activation-recompute-policy.md——一个 prompt，接收你的模型配置（层数、hidden、seq、batch）和可用 GPU 内存，输出一个逐层重算策略（无 / 选择性 / 全 / 卸载）。

练习

1. 验证正确性。跑 model_forward + model_backward（全激活）vs model_forward_checkpointed + model_backward_checkpointed（分段）。参数梯度必须到机器精度相同。

2. 把段大小 k 从 1 扫到 L。画出 FLOP 开销和内存。找出曲线的拐点。

3. 实现选择性检查点：存注意力模块的输入但不存它的中间量。对一个 seq=8192 的 32 层模型，测量它相对全层检查点的 FLOP 开销。

4. 加上卸载。把段输入存到一个模拟的 "CPU 缓冲"（一个单独的列表）。把 "PCIe 带宽" 测量为字节/时间，找出卸载和重算之间的平衡点。

5. 用和不用 torch.utils.checkpoint 给一个真实的 PyTorch transformer 做基准测试。测量内存（经 torch.cuda.max_memory_allocated）和步进时间。

关键术语

术语	人们怎么说	它实际是什么
梯度检查点	"重做前向来省内存"	只存段输入；反向时重算中间量以得到支撑梯度的张量
激活重算	"和检查点一样"	同一技术的 HPC 风味名字
段大小（k）	"每检查点多少层"	其中间量被一起丢弃并重物化的层数
选择性检查点	"Korthikanti 的诀窍"	只重算存储昂贵的激活（注意力 softmax）；保留便宜的
全检查点	"朴素版本"	在每段重算每层的中间量
块检查点	"粗粒度"	检查点整个 transformer 块；最大粒度
FLOP 开销	"计算税"	每步额外 FLOPs = （重算 FLOPs）/（前向 + 反向 FLOPs）；朴素 33%，选择性 5%
激活卸载	"运到 CPU"	在前向->反向间把激活移到 CPU 内存；重算的替代方案
sqrt-L 规则	"经典最优"	对等成本层，最优检查点间距是 sqrt(L) 层
注意力-softmax 体积	"O(L^2) 问题"	L^2 * heads * batch 个 float；在长 context 下主导激活内存

延伸阅读

• 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 -- 标准 API
• Megatron-Core activation recomputation documentation -- 选择性、全和块模式