梯度累积进阶：混合精度与分布式训练

梯度累积本身并不复杂，但一旦跟混合精度训练、分布式训练这些技术结合起来，坑就多了。我刚开始做大规模训练时，就因为这些组合技踩过不少坑。今天咱们就把这些进阶玩法掰开揉碎讲清楚。

梯度累积 + 混合精度训练

混合精度训练（AMP）现在几乎是标配了。它用FP16做计算，用FP32做参数更新，能省一半显存。但跟梯度累积一起用，有个细节特别容易翻车。

先看一个典型的错误写法：

scaler = GradScaler()
accum_steps = 4

for batch in dataloader:
    with autocast():
        loss = model(batch)
        loss = loss / accum_steps  # 手动缩放
    
    scaler.scale(loss).backward()
    
    if (batch_idx + 1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

这段代码看起来没问题，对吧？其实有个隐藏的bug。

问题出在 loss / accum_steps 这一步。混合精度训练时，loss 是 FP16 的。你除以 4，数值精度会丢失。更关键的是，梯度累积的缩放应该在梯度上做，而不是在 loss 上。

我个人习惯的做法是：

scaler = GradScaler()
accum_steps = 4

for batch in dataloader:
    with autocast():
        loss = model(batch)
    
    scaler.scale(loss).backward()
    
    if (batch_idx + 1) % accum_steps == 0:
        # 梯度累积时，需要把梯度除以累积步数
        for param in model.parameters():
            if param.grad is not None:
                param.grad /= accum_steps
        
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

我在项目中遇到过更坑的情况——梯度累积步数很大（比如 32 步），FP16 的梯度累积过程中出现了下溢。解决办法是：用 FP32 累积梯度。

# 为每个参数维护一个 FP32 的梯度累积器
grad_accum = {name: torch.zeros_like(param, dtype=torch.float32) 
              for name, param in model.named_parameters()}

for batch in dataloader:
    with autocast():
        loss = model(batch)
    
    scaler.scale(loss).backward()
    
    # 将 FP16 梯度累积到 FP32 中
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_accum[name] += param.grad.float()
    
    if (batch_idx + 1) % accum_steps == 0:
        # 把累积的 FP32 梯度写回模型
        for name, param in model.named_parameters():
            if param.grad is not None:
                param.grad = grad_accum[name] / accum_steps
        
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
        
        # 清空累积器
        for name in grad_accum:
            grad_accum[name].zero_()

嗯，这里要注意：这个做法会多占一些显存（因为要存 FP32 的梯度累积器），但换来的是数值稳定性。如果你的模型对精度要求高，这个取舍是值得的。

梯度累积 + 分布式训练

分布式训练（DDP）跟梯度累积的组合，是另一个容易出问题的地方。核心问题在于：梯度同步的时机。

先看 DDP 的基本原理：每个 GPU 独立计算梯度，然后在 backward() 时自动同步梯度（通过 all-reduce）。但梯度累积时，我们并不希望每个 micro-batch 都同步——那样通信开销太大了。

正确的做法是：只在累积的最后一步同步梯度。

from torch.nn.parallel import DistributedDataParallel as DDP

model = DDP(model)
accum_steps = 4

for batch_idx, batch in enumerate(dataloader):
    # 关键：只在最后一步同步梯度
    is_last_step = (batch_idx + 1) % accum_steps == 0
    
    with model.no_sync() if not is_last_step else contextlib.nullcontext():
        loss = model(batch)
        loss = loss / accum_steps
        loss.backward()
    
    if is_last_step:
        # 此时梯度已经同步过了
        optimizer.step()
        optimizer.zero_grad()

我曾经犯过一个错误：在分布式训练中，把 loss 除以了 world_size * accum_steps。结果梯度变得特别小，模型根本训不动。后来才意识到，DDP 的 all-reduce 已经做了平均，你只需要除以累积步数就行。

还有一个容易忽略的点：学习率调整。梯度累积相当于增大了 batch size，所以学习率也要相应调整。我一般遵循线性缩放规则：

但要注意，这个规则不是万能的。如果 batch size 太大（比如超过 1024），线性缩放可能失效，需要改用平方根缩放或 warmup 策略。

梯度累积的常见陷阱

做梯度累积这几年，我总结了几条血泪教训：

# 错误：每个 micro-batch 都裁剪
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 正确：只在最后一步裁剪
if is_last_step:
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()

# 在 optimizer 中设置 weight_decay 时，除以累积步数
optimizer = torch.optim.AdamW(model.parameters(), 
                              lr=1e-4, 
                              weight_decay=1e-5 / accum_steps)

知识体系总览

下面这张图总结了梯度累积进阶的核心逻辑：

说实话，梯度累积的这些进阶技巧，刚开始接触时确实容易懵。但只要你理解了核心原理——梯度累积本质上是把 N 个 micro-batch 的梯度加起来，再统一更新参数——那么所有的问题都能迎刃而解。

混合精度训练要小心数值精度，分布式训练要控制同步时机，各种陷阱要提前预防。把这些都搞定了，你就能在显存受限的情况下，训练出跟大 batch size 一样好的模型。

想要更多资料，联系可可