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

参考答案版：包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。

Important统一完成标准与及格线

代码 4 分、Q1～Q3 各 2 分，通过线 8/10。必须解释正确性边界、显存/通信公式中的单位与分片维度；未实际运行的内容只能标记为静态审查。

本课目标

完成后你需要能够：

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

核心概念

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

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

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

mact≈L⋅s⋅b⋅h⋅(34+5⋅nheads⋅sh)bytesm_{\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}

其中参数定义为：模型层数 LL、输入序列长度 ss、微批次大小 bb、隐藏层维度 hh、自注意力头数 nheadsn_{\text{heads}}。

  • 线性增长项 34⋅L⋅s⋅b⋅h34 \cdot L \cdot s \cdot b \cdot h：来自 Transformer 块内的常规张量。包括 Q/K/V 投影输入、注意力输出投影、两层 LayerNorm/RMSNorm 输入、SwiGLU/GeLU 激活函数输入及残差连接。每个张量在 16 位精度下占 22 字节，整层累加约为 34h34h 字节/token。
  • 二次方爆炸项 5⋅L⋅b⋅nheads⋅s25 \cdot L \cdot b \cdot n_{\text{heads}} \cdot s^2：来自自注意力矩阵。计算 QKTQ K^T 会生成形状为 [b,nheads,s,s][b, n_{\text{heads}}, s, s] 的注意力分数矩阵（Attention Score），后续还要保存 Softmax 输出概率矩阵与 Dropout 掩码。由于该部分显存与序列长度呈二次方关系（O(s2)O(s^2)），当序列长度从 2k 扩展至 32k、128k 时，注意力激活值将迅速暴涨并彻底压垮显存！

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

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

flowchart TD
    subgraph FullForward["朴素前向: 全量暂存激活 (显存随层数 L 线性累加)"]
        direction LR
        L1["Layer 1 激活"] --> L2["Layer 2 激活"] --> LN["Layer L 激活 (显存峰值触顶 O(L))"]
    end

    subgraph FullRecompute["全量重计算: 仅留层边界 Checkpoint (显存压缩至 O(1))"]
        direction LR
        CP1["Checkpoint 1"] -.->|"前向丢弃中间激活<br>反向重新计算 (+33% FLOPs)"| RB1["反向重算 L1 激活"]
        CP2["Checkpoint 2"] -.->|"前向丢弃中间激活<br>反向重新计算"| RB2["反向重算 L2 激活"]
    end

    subgraph SelectiveRecompute["选择性重计算 (Selective / FlashAttention 原生)"]
        direction LR
        MLP_Act["MLP GEMM: FLOPs 占 97%, 显存仅线性<br>===> 保留激活"]
        Attn_Act["Attention QK^T: FLOPs 仅 2.7%, 显存占 70% 二次方<br>===> 丢弃并在片上重算 (极高性价比)"]
    end

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

  • 全量重计算（Full Recomputation）：仅在每个 Transformer 层的输入边界保存 Checkpoint，层内部的所有中间激活全部丢弃。在反向传播时，每层都需要重新执行一遍完整前向。
    • 收益：激活显存开销直接压减为原来的约 1/L1/L（几乎只剩单层激活）。
    • 代价：训练计算量增加约 30%∼33%30\% \sim 33\%（因为整个前向计算被执行了两次）。
  • 选择性重计算（Selective Recomputation）：Megatron-LM 与 FlashAttention 提出的重大突破。观察表明：注意力矩阵 QKTQ K^T 占据了 70%70\% 以上的激活显存，但其计算 FLOPs 仅占整层计算的不到 3%3\%！ 相比之下，MLP 层的 GEMM 占据了绝大部分 FLOPs 却只贡献线性显存。因此，选择性重计算只丢弃注意力分数与 Softmax 激活，而保留 MLP 与投影的中间值。
    • 收益：以仅 +2.7%+2.7\% 的轻微计算开销，削减了高达 70%70\% 的激活显存！
    • 与 FlashAttention 的关系：FlashAttention 利用 GPU SRAM 瓦片切分在线计算 Softmax，其反向过程本身就直接在片上 SRAM 重算注意力矩阵，因此使用 FlashAttention 实际上已在算子内部原生实现了选择性重计算。

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

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

flowchart LR
    subgraph MicroBatches["梯度累积 Micro-Batch 调度时序"]
        direction TB
        MB1["Micro-Batch 1 (前向+反向)"] -->|累加梯度| G["梯度缓冲区 (Grad Buffer)"]
        MB2["Micro-Batch 2 (前向+反向)"] -->|累加梯度| G
        MBN["Micro-Batch N (前向+反向)"] -->|累加梯度| G
        G --> Step["除以 N 归一化均值 -> Optimizer Step"]
    end

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

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

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

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

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

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

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

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

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

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

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

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

  • s=512s = 512（短序列）：系数 =34+20=54= 34 + 20 = 54，激活显存 ≈3.6 GB\approx 3.6\text{ GB}（远小于 104 GiB 静态显存，几乎可忽略）。
  • s=2048s = 2048（标准序列）：系数 =34+80=114= 34 + 80 = 114，激活显存 ≈30.6 GB\approx 30.6\text{ GB}（已达到显著量级）。
  • s=4096s = 4096（长序列）：系数 =34+160=194= 34 + 160 = 194，激活显存 ≈104.2 GB\approx 104.2\text{ GB}（二次方项占比反超线性项，且激活显存已与 104 GiB 静态账本持平！）。
  • s=8192s = 8192（超长序列）：系数 =34+320=354= 34 + 320 = 354，激活显存 ≈380.1 GB\approx 380.1\text{ GB}（激活显存成为绝对主导瓶颈，必须开启选择性重计算并结合上下文并行 CP）。

代码填空题

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

Code 
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 状态与缓存分配器预热。

参考答案（仅 answer 分支）

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

Code 
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；反超 112 GB 静态账本大约发生在 seq≈4k 附近（具体数值随模型结构浮动，重点是趋势与量级判断）。

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），公式给出的是组成/下界。

官方主参考

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