详解深度学习中的梯度累积技巧：如何在有限显存下模拟大批量训练

GPU硬件与深度学习训练

在深度学习模型训练中，batch size 的选择直接影响模型的收敛速度和最终精度。研究表明，较大的 batch size 能让梯度估计更加稳定，有助于模型跳出局部最优，同时充分利用 GPU 的并行计算能力。然而，受限于显存容量，很多开发者不得不将 batch size 设得很小，导致训练效果大打折扣。梯度累积（Gradient Accumulation）正是解决这一矛盾的核心技巧——它通过多次前向-反向传播累积梯度，再统一更新参数，从而在不增加显存开销的前提下模拟大批量训练的效果。

深度学习训练流程

Table of Contents

一、为什么需要梯度累积？

假设你有一张 24GB 显存的 RTX 4090，训练一个 ViT-Large 模型时，单张图片前向传播加上反向传播的梯度就已经占用了约 20GB 显存，此时 batch size 最多只能设为 1。但论文中推荐的 batch size 是 256，差距悬殊。梯度累积的思路很直观：把一个大 batch 拆成多个小 mini-batch，依次前向和反向计算，把每个 mini-batch 的梯度累加起来，等累积够了再执行一次参数更新。这样，显存中始终只保存一份模型参数和一份累积梯度的缓冲区，显存占用与小 batch 完全一致。

二、PyTorch 实现梯度累积

在 PyTorch 中实现梯度累积非常简洁，核心在于控制 optimizer.zero_grad() 和 optimizer.step() 的调用时机。以下是一个完整的训练循环示例：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# 模拟数据集
X = torch.randn(1024, 64)
y = torch.randint(0, 10, (1024,))
dataset = TensorDataset(X, y)

# micro_batch_size 是显存能承受的实际 batch size
# accumulation_steps 累积步数，effective_batch = micro_batch_size * accumulation_steps
micro_batch_size = 8
accumulation_steps = 32  # 等效 batch size = 8 * 32 = 256

dataloader = DataLoader(dataset, batch_size=micro_batch_size, shuffle=True)
model = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

model.train()
for epoch in range(10):
    optimizer.zero_grad()  # 每个 epoch 开始时清零一次
    for step, (batch_x, batch_y) in enumerate(dataloader):
        logits = model(batch_x)
        loss = loss_fn(logits, batch_y)
        loss = loss / accumulation_steps  # 关键：对 loss 做平均化处理
        loss.backward()  # 梯度会自动累加到 .grad 属性上

        if (step + 1) % accumulation_steps == 0:
            optimizer.step()  # 累积够了，执行一次参数更新
            optimizer.zero_grad()  # 清零梯度，开始下一轮累积

    # 处理最后一个不足 accumulation_steps 的尾部 batch
    if (step + 1) % accumulation_steps != 0:
        optimizer.step()
        optimizer.zero_grad()

上面代码中有两个关键细节：一是 loss 除以累积步数，因为 PyTorch 的 backward() 会将梯度累加到 .grad 属性上，如果不做缩放，累积后的梯度会是实际值的 N 倍；二是尾部 batch 的处理，当数据集大小不能被 effective batch size 整除时，最后一个不完整的累积周期也需要执行一次参数更新。

三、梯度累积的等价性证明

很多初学者会疑惑：梯度累积真的和直接用大 batch 等价吗？从数学上看，假设损失函数为 L，对参数 θ 求梯度：

大 batch 梯度: g = (1/N) * Σ(∇L_i), i = 1..N
梯度累积:     g = (1/K) * Σ(Σ(∇L_j)/M), 每个 K 内含 M 个样本
等价条件: N = K * M

两者在数值上完全等价，但有一个前提条件：Batch Normalization 的行为会不同。BN 层在每个 mini-batch 内计算均值和方差，小 batch size 下统计量会有较大噪声。解决方案是使用 Group Normalization 或 Layer Normalization 替代，或者在累积期间保持 BN 层的 eval 模式。

四、实战中的注意事项与进阶技巧

在实际工程中使用梯度累积，有几个值得注意的坑和优化方向：

1. 学习率线性缩放：当 effective batch size 从 B 变为 K*B 时，通常需要将学习率也乘以 K（Linear Scaling Rule），以保持每步更新的有效幅度。配合 warmup 使用效果更佳。

from torch.optim.lr_scheduler import LambdaLR

base_lr = 1e-3
scale_factor = accumulation_steps  # 学习率缩放倍数
optimizer = torch.optim.AdamW(model.parameters(), lr=base_lr * scale_factor)

# warmup 策略：前 10% 步数线性升温
warmup_steps = len(dataloader) // accumulation_steps * 10 // 100
def lr_lambda(current_step):
    if current_step < warmup_steps:
        return current_step / warmup_steps
    return 1.0

scheduler = LambdaLR(optimizer, lr_lambda)

2. 梯度裁剪的位置：torch.nn.utils.clip_grad_norm_ 应该在 optimizer.step() 之前调用，即累积完成、更新参数之前。这样才能对累积后的整体梯度做裁剪，避免某个 mini-batch 的异常梯度被放大。

3. 混合精度训练的配合：使用 AMP（Automatic Mixed Precision）时，梯度累积与 GradScaler 的配合需要格外小心。Scaler 的 update() 应在参数更新时调用一次，而不是每个 mini-batch 都调用。

scaler = torch.cuda.amp.GradScaler()

for step, (batch_x, batch_y) in enumerate(dataloader):
    with torch.cuda.amp.autocast():
        logits = model(batch_x)
        loss = loss_fn(logits, batch_y) / accumulation_steps

    scaler.scale(loss).backward()

    if (step + 1) % accumulation_steps == 0:
        scaler.unscale_(optimizer)  # 在裁剪前 unscale
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)  # 内部会检查 inf/nan
        scaler.update()
        optimizer.zero_grad()

代码编程实践

五、与其他显存优化技术的组合

梯度累积并不是孤立的技术，它可以和多种显存优化手段叠加使用，形成更强大的训练方案：

梯度检查点（Gradient Checkpointing）：用时间换空间，在反向传播时重新计算中间激活值，而非全部存储。配合梯度累积，可以在单张消费级显卡上训练更大的模型。PyTorch 通过 torch.utils.checkpoint 实现：

from torch.utils.checkpoint import checkpoint

class LargeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.block1 = nn.Sequential(nn.Linear(512, 512), nn.ReLU())
        self.block2 = nn.Sequential(nn.Linear(512, 512), nn.ReLU())

    def forward(self, x):
        # checkpoint 不保存中间激活，反向时重新计算
        x = checkpoint(self.block1, x, use_reentrant=False)
        x = checkpoint(self.block2, x, use_reentrant=False)
        return x

DeepSpeed ZeRO 与梯度累积：DeepSpeed 的 ZeRO 优化器本身就支持梯度累积，只需在配置文件中设置 gradient_accumulation_steps 即可。ZeRO-Stage2 将梯度和优化器状态分片到多张 GPU 上，再配合梯度累积，能在 4 张 RTX 3090 上训练原本需要 A100 才能跑的模型。

总结来说，梯度累积是一项零成本、高收益的训练技巧。它不引入额外的显存开销，实现简单，且能与混合精度、梯度检查点、分布式训练等技术无缝配合。只要记住三个要点：loss 除以累积步数、更新前做梯度裁剪、注意 BN 层的等价性问题，就能在有限的硬件条件下获得接近大集群的训练效果。