【Debug日志 | DDP + AMP + 梯度累积下有效学习率被放大】
原创
【Debug日志 | DDP + AMP + 梯度累积下有效学习率被放大】
原创
九年义务漏网鲨鱼
修改于 2025-09-08 15:29:51
修改于 2025-09-08 15:29:51
代码可运行
运行总次数：0
代码可运行
训练突然“炸掉”：DDP + AMP + 梯度累积下有效学习率被放大
场景：在多卡 DDP、开启 AMP（混合精度）模式下，同时采用了梯度累积（gradient accumulation） 做大 batch。单卡/小 batch 能学；一上多卡 + 累积就容易发散、loss 爆涨、grad_norm 飙高。在本章中主要阐述bug可能发生的现象以及总结原因。
❓Bug 现象
- 单卡、无累积：loss 正常下降。
- 4 卡 + AMP +
accum_steps=8：几十个 step 内 loss 爆炸，grad_norm经常 > 1e3；把 LR 降低 8× 似乎能“救回一点”，但仍然不稳。 - 关闭 AMP 后只稍有改善，仍不稳。
📽️ 场景复现
典型错误：没有对 loss 除以
accum_steps；每个小步都在调用scheduler.step()；clip_grad_norm_在unscale_之前或在每小步都裁剪。
import os, torch, torch.nn as nn, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def train_one_epoch(model, loader, optimizer, scheduler, scaler, accum_steps=8, max_norm=1.0):
model.train()
for step, (x, y) in enumerate(loader):
x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
with torch.cuda.amp.autocast(True):
logits = model(x)
loss = nn.functional.cross_entropy(logits, y) # ❌ 未缩放
optimizer.zero_grad(set_to_none=True) if step % accum_steps == 0 else None
scaler.scale(loss).backward() # ❌ 累积但没 /accum_steps
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) # ❌ 未 unscale_ 就裁剪
if (step + 1) % accum_steps == 0:
scaler.step(optimizer)
scaler.update()
scheduler.step() # ❌ 每小更新都在 step 调度- 可能原因
- 有效 batch =
per_gpu_bs × world_size × accum_steps，但loss 没除以accum_steps，等价于把梯度放大accum_steps倍； - 若还按每小步调用
scheduler.step()，学习率在一个 epoch 内被“走太快”；
Debug 过程
1️⃣ 打印“有效学习率”与梯度范数
eff_lr = optimizer.param_groups[0]["lr"] * (accum_steps)
print(f"lr={optimizer.param_groups[0]['lr']:.2e}, eff_lr≈{eff_lr:.2e}")- 现象：把理论上的
lr乘上accum_steps后，与发散时的感觉吻合（太大）。
同时在关键层上挂grad hook：
def watch_grad(module, name):
def _hook(grad):
print(name, grad.norm().item())
return _hook
model.layer4.register_full_backward_hook(lambda m, gin, gout: print('g=', gout[0].norm()))- 现象：每
accum_steps内的每个小步梯度范数都很高，最终一步step()前已严重超标。
2️⃣ 分离“优化步”和“调度步”
- 人为把
scheduler.step()改到优化器更新后，并且只在完成一次累积后才步进，震荡明显减弱。
3️⃣ AMP 裁剪顺序验证
- 在
scaler.unscale_(optimizer)之后再裁剪，grad_norm日志与直觉一致。
✅ 修复方案（稳定模板）
# fixed_train.py —— 可直接用
import os, torch, torch.nn as nn, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def train_one_epoch(model, loader, optimizer, scheduler, scaler,
accum_steps=8, max_norm=1.0):
model.train()
optimizer.zero_grad(set_to_none=True)
for step, (x, y) in enumerate(loader):
x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
with torch.cuda.amp.autocast(True):
logits = model(x)
loss = nn.functional.cross_entropy(logits, y)
loss = loss / accum_steps
scaler.scale(loss).backward()
if (step + 1) % accum_steps == 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer) # 自动跳过 inf/NaN
scaler.update()
optimizer.zero_grad(set_to_none=True)
if scheduler is not None:
scheduler.step()loss /= accum_steps；- 只在
(step+1) % accum_steps == 0时：unscale_ → clip → step → update → scheduler.step()； - 其他小步只做
scaled backward，不要调度、不要裁剪、更不要step()； zero_grad建议放在优化步后，且用set_to_none=True。
调度器与“每步/每轮”配套
- 每步调度（如 CosineAnnealingLR T_max=总优化步数）：确保 T_max = (num_batches / accum_steps) × epochs；
- 每轮调度（StepLR/ReduceLROnPlateau）：只在 epoch 末
step()，与累积无关。
示例（Cosine 每步调度）：
total_updates = (len(train_loader) // accum_steps) * num_epochs
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_updates)结果
- 4 卡、
accum_steps=8、Cosine 每步调度： - 发散消失，
grad_norm稳定在 0.5–3 范围； - 同等有效 batch下，与“无累积、单大 batch”的收敛曲线对齐（±1e-3）；
- AMP 加速正常（1.6×–2.1×）。
梯度累积本质是把多次小反传合成一次优化步。只要记住两点：
- 把损失按累积分摊；
- 一切与“优化步”绑定的动作（裁剪、step、调度、zero_grad）都只在最后一步做， 多卡 + AMP + 累积的组合就会稳定而高效。你可以直接把上面的 fixed 模板拷进项目，把“自检脚本”加到训练循环里，后续换配置也不容易踩坑。
- 没有按累积步数缩放 loss、调度器步进时机错，导致有效学习率被放大 N 倍。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
评论
登录后参与评论
推荐阅读
目录
腾讯云开发者
