📚 课程目录第3章梯度累积原理

3、梯度累积原理：小批量拆分、累积步数、等效批量大小、内存与计算权衡

在单卡显存无法容纳大 batch 时，梯度累积（Gradient Accumulation）是最常用的“以时间换空间”手段。它的核心思想非常朴素：把一个大 batch 拆成若干个 micro-batch，逐个前向与反向，把每个 micro-batch 产生的梯度累加到同一组参数梯度缓冲区中，等累积到指定步数后再统一执行一次优化器更新。这样，模型参数只更新一次，但参与计算的样本数却等效于一个大 batch。

梯度累积不改变模型结构，也不改变优化器算法，它改变的是“梯度计算”与“参数更新”之间的时序关系：先累积，后更新。

一、小批量拆分与累积步数

设全局等效批量大小为 B_global，单次前向能放入显存的 micro-batch 大小为 B_micro，则累积步数 K 满足：

K = B_global / B_micro

例如：
B_global = 1024
B_micro  = 64
K = 1024 / 64 = 16

训练时，数据加载器每次只取 B_micro 条样本，连续执行 K 次前向和反向。每次反向结束后，梯度被累加到 param.grad 中，而不是立即调用 optimizer.step()。当第 K 次反向完成后，才执行一次参数更新，并清空梯度。伪代码如下：

for step, micro_batch in enumerate(loader):
    outputs = model(micro_batch)
    loss = criterion(outputs, targets) / K
    loss.backward()

    if (step + 1) % K == 0:
        optimizer.step()
        optimizer.zero_grad()

注意 loss / K 这一步非常关键。因为 PyTorch 的 backward() 默认是累加梯度，如果不对 loss 做缩放，累积后的梯度将是真实梯度的 K 倍，等效于把学习率放大了 K 倍，容易导致训练发散。

二、等效批量大小与梯度关系

从数学上看，若每个 micro-batch 的损失为 L_i，则全局 batch 的损失为：

L_global = (1/K) * Σ L_i

对参数 θ 求导：

∂L_global/∂θ = (1/K) * Σ ∂L_i/∂θ

这正是梯度累积所做的事情：把 K 个 micro-batch 的梯度求平均。因此，在数值上，梯度累积后的梯度与直接使用 B_global 的大 batch 梯度是等价的（忽略 BatchNorm 等对 batch 维度敏感的层）。

三、内存与计算权衡

梯度累积最直接的好处是显存占用显著降低。假设模型参数量为 P，优化器状态为 S，激活值为 A(B)，则显存大致为：

显存 ≈ P + S + A(B_micro) + 梯度缓冲区

其中：
P：模型参数
S：优化器状态（如 Adam 的 m、v）
A(B_micro)：micro-batch 的激活值
梯度缓冲区：与 P 同量级，但只有一份

如果直接使用 B_global，激活值 A(B_global) 会随 batch 线性增长，往往成为显存瓶颈。而梯度累积把激活值限制在 A(B_micro)，代价是前向和反向需要执行 K 次，计算时间近似增加 K 倍。因此，梯度累积本质上是“用计算时间换显存空间”。

维度	直接大 batch	梯度累积
等效批量	B_global	B_micro × K
激活显存	随 B_global 线性增长	仅与 B_micro 相关
计算次数	1 次前向 + 1 次反向	K 次前向 + K 次反向
参数更新频率	每个 batch 更新一次	每 K 个 micro-batch 更新一次
适用场景	显存充足	显存受限、需大 batch 稳定训练

四、案例与常见误区

假设我们在单卡 24GB 显存上训练一个 7B 模型，直接使用 batch size 512 会 OOM。此时可以设置 B_micro = 32，K = 16，等效批量仍为 512。训练脚本中只需把 loss 除以 16，并在每 16 步后执行一次优化器更新。这样显存峰值从“无法加载”降到“可稳定运行”，但单步耗时约增加 16 倍，整体吞吐下降。

常见误区一：忘记对 loss 做除法。结果梯度被放大 K 倍，等效于学习率乘以 K，训练很快发散。
常见误区二：在累积过程中调用 optimizer.zero_grad()。这会把之前累积的梯度清空，导致只有最后一个 micro-batch 生效。
常见误区三：认为梯度累积可以完全等价于大 batch。对于 BatchNorm、LayerNorm 等依赖 batch 统计量的层，micro-batch 的统计量与全局 batch 不同，可能带来精度差异。

实操建议：在梯度累积时，优先使用 loss / K 而不是 loss / B_global，因为 K 是累积步数，语义更清晰；同时配合梯度裁剪（gradient clipping）时，应在 optimizer.step() 之前、累积完成后进行裁剪，避免每个 micro-batch 单独裁剪导致尺度不一致。

总结来说，梯度累积通过“小批量拆分 + 梯度累加 + 延迟更新”三个动作，把大 batch 的训练需求映射到小显存设备上。理解 B_global = B_micro × K 这一核心关系，并正确处理 loss 缩放与更新时机，就能在显存与计算之间取得可控的平衡。