第 2 课：激活、重计算与梯度累积

• 对应官方章节：Memory usage in transformers（激活部分）/ Activation recomputation / Gradient accumulation
• 前置：第 1 课（静态显存账本）
• 状态：未开始

本课目标

完成后你需要能够：

• 使用激活显存公式估算并说明其随 batch size 与序列长度的增长方式。
• 解释 full / selective 重计算的内存-计算取舍，并说明与 FlashAttention 的关系。
• 用代码实现正确的梯度累积（含平均、清零时机），理解其与 DP 的关系（第 4 课展开）。
• 区分”峰值显存”与”平均/静态显存”，并知道如何测量峰值。

核心概念

1. 动态激活值的本质与二次方爆炸

在前向传播过程中，神经网络中的非线性算子、LayerNorm 与多头自注意力机制需要保存中间输入，供反向传播根据微积分链式法则求导（例如求 $\frac{\partial L}{\partial W}$ 需要用到输入 $X$：$\frac{\partial L}{\partial W} = X^T \cdot \frac{\partial L}{\partial Y}$）。

在大语言模型（Transformer）中，混合精度（16-bit）下 $L$ 层的总激活显存由两部分构成：

$$m_{\text{act}} \approx L \cdot s \cdot b \cdot h \cdot \left(34 + 5 \cdot \frac{n_{\text{heads}} \cdot s}{h}\right) \quad \text{bytes}$$

其中参数定义为：模型层数 $L$、输入序列长度 $s$、微批次大小 $b$、隐藏层维度 $h$、自注意力头数 $n_{\text{heads}}$。

• 线性增长项 $34 \cdot L \cdot s \cdot b \cdot h$：来自 Transformer 块内的常规张量。包括 Q/K/V 投影输入、注意力输出投影输入、两层 LayerNorm 输入、MLP 的 GeLU 输入及 dropout 掩码。16 位张量每元素占 $2$ 字节、dropout 掩码每元素占 $1$ 字节，整层累加约为 $34h$ 字节/token。
  • 注意：系数 34 按 Korthikanti 等对 GPT 块的推导得出（MLP 中间维度 $4h$、GeLU、LayerNorm、带 dropout）。LLaMA 类模型使用 SwiGLU、RMSNorm 且通常不用 dropout，系数会不同，这里只用于量级估算。
• 二次方爆炸项 $5 \cdot L \cdot b \cdot n_{\text{heads}} \cdot s^2$：来自自注意力矩阵。计算 $Q K^T$ 会生成形状为 $[b, n_{\text{heads}}, s, s]$ 的注意力分数矩阵（Attention Score），后续还要保存 Softmax 输出概率矩阵与 Dropout 掩码。由于该部分显存与序列长度呈二次方关系（$O(s^2)$），当序列长度从 2k 扩展至 32k、128k 时，注意力激活值将迅速暴涨并压垮显存。

2. 激活重计算（Activation Checkpointing / Rematerialization）

激活重计算的核心思想是以算力换显存：在前向传播中主动丢弃部分或全部中间激活，反向传播到达该层时，再从最近保存的检查点（Checkpoint）重新计算前向（如图 2-1 所示）：

朴素前向: 全量暂存激活 (显存随层数 L 线性累加)

Layer 1 激活

Layer 2 激活

Layer L 激活 (显存峰值触顶 O(L))

全量重计算: 仅留层边界 Checkpoint (检查点 2sbh·L 仍为 O(L), 另加单层重算激活)

Checkpoint 1

反向重算 L1 激活

Checkpoint 2

反向重算 L2 激活

选择性重计算 (Selective / FlashAttention 原生)

MLP / 投影 GEMM: 承担绝大部分 FLOPs, 显存仅线性
===> 保留激活

QK^T / Softmax / Dropout 的 s² 激活
显存占比大、重算便宜 → 反向时重算

前向丢弃中间激活
反向重新计算 (+33% FLOPs)前向丢弃中间激活
反向重新计算

图 2-1：全量重计算 vs 选择性重计算显存开销与计算权衡架构

• 全量重计算（Full Recomputation）：仅在每个 Transformer 层的输入边界保存 Checkpoint，层内部的所有中间激活全部丢弃。在反向传播时，每层都需要重新执行一遍完整前向。
  • 收益：激活显存降为层边界检查点（每层 $2sbh$ 字节，共 $2sbh \cdot L$，仍随 $L$ 线性增长）加上反向时正在重算的单层激活，远小于全量保存。
    
  • 代价：训练计算量增加约 $30\% \sim 33\%$（因为整个前向计算被执行了两次）。
• 选择性重计算（Selective Recomputation）：由 Korthikanti 等（Megatron-LM）提出。观察表明：注意力中 $QK^T$、Softmax、Dropout 产生的 $s^2$ 项张量占激活显存的大头，而重算它们所需的 FLOPs 很少；MLP 与投影 GEMM 承担绝大部分 FLOPs，却只贡献线性显存。因此，选择性重计算只丢弃注意力分数与 Softmax 相关激活，而保留 MLP 与投影的中间值。
  • 收益：论文在 GPT-3 175B 上报告，以约 $+2.7\%$ 的 FLOPs 开销（这是重算带来的额外开销，不是注意力在整层 FLOPs 中的占比）削减约 $70\%$ 的激活显存。
  • 与 FlashAttention 的关系：FlashAttention 利用 GPU SRAM 瓦片切分在线计算 Softmax，其反向过程本身就直接在片上 SRAM 重算注意力矩阵，因此使用 FlashAttention 实际上已在算子内部原生实现了选择性重计算。

3. 梯度累积（Gradient Accumulation）的工作机制

在单卡显存受限、无法容纳目标全局 Batch Size 时，梯度累积将大批次划分为 $N_{\text{acc}}$ 个微批次（Micro-Batch），如图 2-2 所示：

梯度累积 Micro-Batch 调度时序

Micro-Batch 1 (前向+反向)

Micro-Batch 2 (前向+反向)

Micro-Batch N (前向+反向)

梯度缓冲区 (Grad Buffer)

除以 N 归一化均值 -> Optimizer Step

累加梯度累加梯度累加梯度

图 2-2：梯度累积跨微批次累加与优化器单步平均执行流

$$\text{全局批大小 (GBS)} = \text{微批大小 (MBS)} \times \text{累积步数 (grad\_acc)} \times \text{数据并行度 (DP)}$$

• 执行流程：每个 Micro-Batch 独立执行前向与反向，并将求得的梯度直接累加到一个常驻的梯度缓冲区（Grad Buffer）中；在累积满 $N_{\text{acc}}$ 步后，执行平均化归一（除以 $N_{\text{acc}}$），随后触发一次优化器更新，最后清空梯度缓冲区。
• 显存与性能权衡：
  • 显存优势：任意时刻显存中仅需保留 1 个 Micro-Batch 的动态激活，激活峰值与全局 Batch 大小完全解耦。
    
  • 性能代价：多个 Micro-Batch 在单卡上串行执行，增加了显存中梯度缓冲区的驻留开销，且引入了微批次循环的调度开销。

4. 显存峰值与“第一步成功后续 OOM”陷阱

在系统调优中，经常出现“第一步训练正常通过，第二步突然报 CUDA OOM”的现象，其物理根因在于：

1. 优化器状态的延迟初始化：在标准 PyTorch 训练循环中，Adam 的动量 $m$ 和方差 $v$（共 $8N$ 字节）通常在首次调用 optimizer.step() 时才被惰性分配（Lazy Allocation）到显存中。第一步前向反向时这 $8N$ 字节尚未存在，第二步前向时静态底噪已暴增，从而触发 OOM。

2. Caching Allocator 预热与显存碎片：首个 Step 会触发大量的显存池预分配，随着多步迭代中张量形状变动与张量生命周期交错，显存池碎片逐渐增多。

3. 测量标准：必须使用 torch.cuda.max_memory_allocated() 测量真实物理张量峰值，而非 nvidia-smi 报告的驱动级预留显存（Reserved Memory）。

具体演示：序列长度增加下的激活显存爆炸

以 LLaMA-7B 典型超参数（$L=32, h=4096, n_{\text{heads}}=32, b=1$）为例，计算混合精度下无重计算的单样本激活显存：

$$\text{系数} = 34 + 5 \cdot \frac{32 \cdot s}{4096} = 34 + \frac{5s}{128}$$

$$\text{总激活显存} \approx 32 \times s \times 1 \times 4096 \times \left(34 + \frac{5s}{128}\right) \text{ bytes}$$

• $s = 512$（短序列）：系数 $= 34 + 20 = 54$，激活显存 $\approx 3.6\text{ GB}$（远小于 104 GiB 静态显存，几乎可忽略）。
• $s = 2048$（标准序列）：系数 $= 34 + 80 = 114$，激活显存 $\approx 30.6\text{ GB}$（已达到显著量级）。
• $s = 4096$（长序列）：系数 $= 34 + 160 = 194$，激活显存 $\approx 104.2\text{ GB}$（$\approx 97.0\text{ GiB}$），已接近 112 GB（104.3 GiB）的静态账本；按本公式，激活在 $s \approx 4260$ 时超过静态账本。二次方项 $5s/128$ 早在 $s > 870$ 时就已超过线性项 34。
• $s = 8192$（超长序列）：系数 $= 34 + 320 = 354$，激活显存 $\approx 380.1\text{ GB}$（激活显存成为绝对主导瓶颈，必须开启选择性重计算并结合上下文并行 CP）。

代码填空题

补全以下程序。注意：不要在估算函数里写死常数总和，所有分量都要显式出现。

import torch

def estimate_activation_memory(
    num_layers: int,
    seq_len: int,
    micro_batch: int,
    hidden: int,
    num_heads: int,
    recompute: str = "none",
) -> float:
    """
    估算混合精度下激活显存（bytes），并可选应用重计算。

    公式（Playbook）：
        m_act ≈ L * seq * bs * h * (34 + 5 * n_heads * seq / h)

    recompute 取值：
        "none"     : 不重计算，保存全部中间激活
        "full"     : 只在层边界保存，约为 none 的 1/num_layers（近似）
        "selective": 丢弃注意力相关激活，约为 none 的 30%（GPT-3 数据，
                     70% 削减、2.7% 计算开销）
    """
    # 每层激活的"每 token 每隐藏单元"系数（含注意力分数二次项）
    per_token_coeff = ______  # 提示：把公式按 L*seq*bs 提出公因子

    base = num_layers * seq_len * micro_batch * hidden * per_token_coeff

    if recompute == "none":
        return base
    elif recompute == "full":
        # full 重计算：只在层边界保留一份激活，约等于每层中间激活全部丢弃
        return ______
    elif recompute == "selective":
        # selective：丢弃注意力激活（约 70%），保留其余
        return ______
    else:
        raise ValueError(f"unknown recompute mode: {recompute}")

def run_gradient_accumulation(
    compute_grads,       # 函数：输入 micro-batch 编号，返回梯度 dict {name: tensor}
    num_microbatches: int,
    grad_accum: int,
    optimizer_step,      # 函数：输入平均后的梯度 dict
) -> None:
    """
    模拟梯度累积训练循环。

    正确性要求：
    1. 每 grad_accum 个 micro-batch 做一次 optimizer step；
    2. optimizer step 使用的是【平均】梯度（除以 grad_accum），
       保证结果与 grad_accum 无关；
    3. step 后必须清零累积缓冲区（否则梯度会被重复计入下一轮）。
    """
    accumulated = None          # 跨 micro-batch 持久的梯度累积缓冲区
    assert num_microbatches % grad_accum == 0

    for mb in range(num_microbatches):
        grads = compute_grads(mb)

        if accumulated is None:
            # 首个 micro-batch：深拷贝作为累积起点
            accumulated = {k: v.clone() for k, v in grads.items()}
        else:
            # 后续 micro-batch：累加
            for k in accumulated:
                accumulated[k] = ______

        # 每攒够 grad_accum 个 micro-batch，做一次优化
        if ______:
            for k in accumulated:
                # 平均，使更新与 grad_accum 无关
                accumulated[k] = ______
            optimizer_step(accumulated)
            # 清零：绝不能漏
            accumulated = ______

def demo_grad_accumulation():
    """用小例子验证：grad_accum=2 时的平均梯度 == 整批一次算出的梯度。"""
    torch.manual_seed(0)
    x = torch.randn(8, 4)
    w = torch.randn(4, 1, requires_grad=True)

    def grads_of(idx):
        w.grad = None
        y = x[idx] @ w
        loss = y.square().mean()
        loss.backward()
        return w.grad.clone()

    # 参照：一次算完整 batch
    full = grads_of(slice(0, 8))

    # 梯度累积：两个 micro-batch 各自 backward 后取平均
    g1 = grads_of(slice(0, 4))
    g2 = grads_of(slice(4, 8))
    accumulated = (g1 + g2) / 2

    print("full:", full.flatten()[:3])
    print("accumulated:", accumulated.flatten()[:3])
    assert torch.allclose(full, accumulated)
    print("梯度累积（平均后）与整批一次 backward 一致 ✓")

if __name__ == "__main__":
    # 验证激活估算的趋势：seq 翻倍时激活近似翻倍以上（二次项），而静态账本不变
    for seq in (512, 1024, 2048, 4096):
        act = estimate_activation_memory(32, seq, 1, 4096, 32)
        print(f"seq={seq:5d}  activations={act / 1e9:8.2f} GB")
    demo_grad_accumulation()

三个问答题

Q1

激活显存公式 L·seq·bs·h·(34 + 5·n_heads·seq/h) 中，两项分别来自哪里？为什么说激活随序列长度近似二次增长、而第 1 课的静态账本完全不受序列长度影响？在什么序列长度下激活会反超静态账本（用 7B、L=32、h=4096、n_heads=32、bs=1 粗算）？

Q2

Full 与 selective 重计算各存什么、丢什么、代价分别是什么？为什么说”FlashAttention 用户已经默认在使用 selective 重计算”？重计算增加了 FLOPs，为什么反而可能让训练更快？

Q3

用 torch.cuda.max_memory_allocated() 测出的”峰值”与 nvidia-smi 看到的”当前占用”有什么区别？为什么第一轮训练的显存 profile 与后续轮不同（至少给出两个原因）？这对”第一轮不 OOM 就认为配置安全”的推断意味着什么？

检查与通过标准

总分 10 分：代码正确 4 分（激活公式、重计算系数、梯度累积平均/清零/触发时机各占一部分）、三题各 2 分、通过线 8 分。

一票否决项：

• 梯度累积不平均（直接求和后 step）且说不出理由。
• 认为重计算”必然变慢”或”必然变快”，不讨论访存-计算取舍。
• 把激活公式中的二次项忽略，声称激活随 seq 线性增长。
• 用”第一轮成功”直接证明配置安全，未提及 optimizer 状态与缓存分配器预热。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对；核对后必须解释关键公式，并改一个规模重新计算。

import torch

def estimate_activation_memory(
    num_layers: int,
    seq_len: int,
    micro_batch: int,
    hidden: int,
    num_heads: int,
    recompute: str = "none",
) -> float:
    """
    估算混合精度下激活显存（bytes），并可选应用重计算。

    公式（Playbook）：
        m_act ≈ L * seq * bs * h * (34 + 5 * n_heads * seq / h)

    recompute 取值：
        "none"     : 不重计算，保存全部中间激活
        "full"     : 只在层边界保存，约为 none 的 1/num_layers（近似）
        "selective": 丢弃注意力相关激活，约为 none 的 30%（GPT-3 数据，
                     70% 削减、2.7% 计算开销）
    """
    # 每层激活的"每 token 每隐藏单元"系数（含注意力分数二次项）
    per_token_coeff = 34 + 5 * num_heads * seq_len / hidden  # 提示：把公式按 L*seq*bs 提出公因子

    base = num_layers * seq_len * micro_batch * hidden * per_token_coeff

    if recompute == "none":
        return base
    elif recompute == "full":
        # full 重计算：只在层边界保留一份激活，约等于每层中间激活全部丢弃
        return base / num_layers
    elif recompute == "selective":
        # selective：丢弃注意力激活（约 70%），保留其余
        return base * 0.3
    else:
        raise ValueError(f"unknown recompute mode: {recompute}")

def run_gradient_accumulation(
    compute_grads,       # 函数：输入 micro-batch 编号，返回梯度 dict {name: tensor}
    num_microbatches: int,
    grad_accum: int,
    optimizer_step,      # 函数：输入平均后的梯度 dict
) -> None:
    """
    模拟梯度累积训练循环。

    正确性要求：
    1. 每 grad_accum 个 micro-batch 做一次 optimizer step；
    2. optimizer step 使用的是【平均】梯度（除以 grad_accum），
       保证结果与 grad_accum 无关；
    3. step 后必须清零累积缓冲区（否则梯度会被重复计入下一轮）。
    """
    accumulated = None          # 跨 micro-batch 持久的梯度累积缓冲区
    assert num_microbatches % grad_accum == 0

    for mb in range(num_microbatches):
        grads = compute_grads(mb)

        if accumulated is None:
            # 首个 micro-batch：深拷贝作为累积起点
            accumulated = {k: v.clone() for k, v in grads.items()}
        else:
            # 后续 micro-batch：累加
            for k in accumulated:
                accumulated[k] = accumulated[k] + grads[k]

        # 每攒够 grad_accum 个 micro-batch，做一次优化
        if (mb + 1) % grad_accum == 0:
            for k in accumulated:
                # 平均，使更新与 grad_accum 无关
                accumulated[k] = accumulated[k] / grad_accum
            optimizer_step(accumulated)
            # 清零：绝不能漏
            accumulated = None

def demo_grad_accumulation():
    """用小例子验证：grad_accum=2 时的平均梯度 == 整批一次算出的梯度。"""
    torch.manual_seed(0)
    x = torch.randn(8, 4)
    w = torch.randn(4, 1, requires_grad=True)

    def grads_of(idx):
        w.grad = None
        y = x[idx] @ w
        loss = y.square().mean()
        loss.backward()
        return w.grad.clone()

    # 参照：一次算完整 batch
    full = grads_of(slice(0, 8))

    # 梯度累积：两个 micro-batch 各自 backward 后取平均
    g1 = grads_of(slice(0, 4))
    g2 = grads_of(slice(4, 8))
    accumulated = (g1 + g2) / 2

    print("full:", full.flatten()[:3])
    print("accumulated:", accumulated.flatten()[:3])
    assert torch.allclose(full, accumulated)
    print("梯度累积（平均后）与整批一次 backward 一致 ✓")

if __name__ == "__main__":
    # 验证激活估算的趋势：seq 翻倍时激活近似翻倍以上（二次项），而静态账本不变
    for seq in (512, 1024, 2048, 4096):
        act = estimate_activation_memory(32, seq, 1, 4096, 32)
        print(f"seq={seq:5d}  activations={act / 1e9:8.2f} GB")
    demo_grad_accumulation()

补全后的代码（关键填空处）

def estimate_activation_memory(num_layers, seq_len, micro_batch, hidden, num_heads, recompute="none"):
    per_token_coeff = 34 + 5 * num_heads * seq_len / hidden   # 提出 L*seq*bs 公因子
    base = num_layers * seq_len * micro_batch * hidden * per_token_coeff
    if recompute == "none":
        return base
    elif recompute == "full":
        return base / num_layers          # 只在层边界保存，近似 1/L
    elif recompute == "selective":
        return base * 0.3                 # 丢弃注意力激活，保留 30%（GPT-3 数据）
    else:
        raise ValueError(...)

梯度累积：

else:
            for k in accumulated:
                accumulated[k] = accumulated[k] + grads[k]     # 或 +=

        if (mb + 1) % grad_accum == 0:
            for k in accumulated:
                accumulated[k] = accumulated[k] / grad_accum   # 平均
            optimizer_step(accumulated)
            accumulated = None                                 # 清零

三个问答题要点

Q1

• 第一项 34·L·seq·bs·h：逐层中间激活（线性层输入/输出、LN 输出等）的线性部分；
• 第二项 5·n_heads·seq²·bs：注意力分数矩阵 S=QKᵀ（每头 seq×seq）的贡献，随 seq 二次增长。
• 静态账本（参数/梯度/优化器状态）与 batch、seq 无关；激活随 bs 线性、随 seq 近似二次。
• 粗算：7B（L=32, h=4096, 头=32, bs=1）在 seq≈2048 时激活 ≈ 32·2048·1·4096·(34+5·32·2048/4096) ≈ 32·2048·4096·(34+80) ≈ 30.6 GB；seq=4096 时 ≈ 32·4096·4096·(34+160) ≈ 104 GB；按此公式在 seq≈4260 时反超 112 GB 静态账本（具体数值随模型结构浮动，重点是趋势与量级判断）。

Q2

• Full：只存层边界激活（每层重算一遍前向），内存最小、计算 +30~40%；
• Selective：丢弃”占内存大且重算便宜”的激活（注意力类），GPT-3 上激活 −70%、计算 +2.7%；
• FlashAttention 反向本就重算注意力分数（不存 S），所以用 FA 等于默认 selective；
• 重计算”多算反而更快”的原因：GPU 上 HBM 访存比计算更贵，少写读 HBM 可以抵消多余 FLOPs；FLOPs 不是唯一成本。

Q3

• max_memory_allocated 是分配器视角的峰值（本进程逻辑占用），nvidia-smi 是设备实时占用（含 CUDA context、其他进程、分配器预留）；两者量级和语义都不同；
• 第一轮不同：a) caching allocator 预热/预分配；b) optimizer 状态在第一次 step 后才完整建立；c) 首次 kernel/context 初始化；
• 结论：不能把”第一轮不 OOM”当作配置安全证明——必须看稳态多轮的峰值曲线（max_memory_allocated），并预留碎片/预留缓冲的余量。

通过要点

• 梯度累积必须平均（除以 grad_acc）且 step 后清零；触发时机是”第 (mb+1) 个 micro-batch 结束时”。
• 重计算的本质是访存-计算权衡，方向因硬件而定。
• 峰值显存需实测（max_memory_allocated），公式给出的是组成/下界。

官方主参考

• Ultra-Scale Playbook
• PyTorch distributed documentation

← 上一课
第 1 课：训练流程与显存账本 下一课 →
第 3 课：分布式通信原语