如何调试梯度累积训练任务

梯度累积（Gradient Accumulation）是一种在显存受限时“用时间换空间”的常用技巧：把大 batch 拆成多个小 step，多次前向/反向后只更新一次参数。调试这类任务时，核心目标是确认“累积版”和“等效大 batch 版”在数学与行为上一致，并排查实现细节导致的偏差。

下面按“从理论到实操”的顺序说明。

一、先明确梯度累积的标准写法

以 PyTorch 为例，正确逻辑一般是：

accum_steps = 4
optimizer.zero_grad()
for i, (x, y) in enumerate(dataloader):
pred = model(x)
loss = criterion(pred, y) / accum_steps   # 关键：除以累积步数
loss.backward()
if (i + 1) % accum_steps == 0:
optimizer.step()
optimizer.zero_grad()

注意三点：

1. loss 要除以 accum_steps（除非你手动缩放 lr）
2. zero_grad 只在更新前调用
3. 每个 micro-batch 的 shuffle 方式要与全 batch 一致

二、调试 checklist（最重要）

1. 数值一致性测试（强烈推荐）

用极小数据，对比：

• 方案 A：batch_size = 32，一步更新
• 方案 B：batch_size = 8，累积 4 步

断言：

• 参数更新后的值完全一致（或浮点误差级一致）
• loss 总和一致（B 的 4 个 loss 加起来 ≈ A 的 loss × 4）

如果不一致，优先检查：

• loss 是否忘了 / accum_steps
• DataLoader 的 shuffle / sampler 是否不同
• dropout / batchnorm 在 micro-batch 上行为不同（见下）

2. BatchNorm / LayerNorm 问题

• BatchNorm：在 micro-batch 上统计的是小 batch 均值/方差，与大 batch 不同

→ 调试时可用 model.eval() 或改用 SyncBatchNorm / GroupNorm

• LayerNorm：通常无影响

验证方法：

• 固定 seed
• 对比 BN 的 running_mean 更新轨迹

3. 学习率 & 优化器状态

梯度累积不改变学习率语义，但容易误用：

• 不要因为“看起来 step 少”就调大 lr
• Adam 的 momentum / variance 是按 step 更新的，不是按 sample

调试技巧：

• 打印 param.grad 的 L2 norm
• 对比累积前后 grad 是否线性叠加

4. 梯度是否为 None（最常见坑）

检查：

• 某些参数没参与计算图
• require_grad=False 被误设
• 多卡 + 累积时，某 rank 没反向

调试代码：

for name, p in model.named_parameters():
if p.grad is None:
print("No grad:", name)

5. 分布式下的梯度累积

在 DDP / FSDP 中：

• backward 仍会做 all-reduce
• 累积的是已同步的梯度

常见错误：

• 在 micro-step 内调用 no_sync() 以外的方式同步
• 忘记 model.no_sync() 导致通信爆炸

正确示例（DDP）：

for i, (x, y) in enumerate(loader):
if i % accum_steps != 0:
with model.no_sync():
loss.backward()
else:
loss.backward()

三、常用调试工具

• 梯度监控
• torch.autograd.detect_anomaly()
• wandb / tensorboard 画 grad norm
• 数值对比
• 固定 seed + 小数据重放
• 日志
• 每个 micro-step 的 loss
• 更新前后的 param diff

四、典型 bug 速查表

如果你愿意，可以把具体框架（PyTorch / TF / JAX）和代码片段发我，我可以直接帮你定位问题。