分布式训练中梯度归一化技巧踩坑记录
最近在优化一个分布式大模型训练项目时，遇到了一个令人头疼的问题：不同rank之间的梯度幅度差异巨大，导致训练不稳定甚至发散。经过一番排查和尝试，总结出几个梯度归一化的实用技巧。
问题背景
在使用PyTorch DDP训练7B参数模型时，发现rank 0的梯度均值约为1e-3，而其他rank的梯度均值却达到了1e-1级别。这种差异导致优化器更新步长严重失衡。
实战技巧
1. 梯度裁剪+归一化组合拳
# 在optimizer.step()前添加
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# 然后归一化处理
for param in model.parameters():
if param.grad is not None:
param.grad /= torch.norm(param.grad) + 1e-8
2. Rank间梯度同步归一化
# 在所有reduce操作后进行
grads = [param.grad for param in model.parameters() if param.grad is not None]
all_grads = torch.cat([g.flatten() for g in grads])
global_norm = torch.norm(all_grads)
# 按全局范数归一化
for param in model.parameters():
if param.grad is not None:
param.grad /= global_norm + 1e-8
3. 动态学习率调整 根据梯度归一化后的幅度动态调整lr，避免训练过快或过慢。
这些方法在实际项目中确实有效，但要注意：梯度过小会降低收敛速度，过大则可能导致训练不稳定。建议在训练过程中持续监控各rank梯度统计信息，及时调整参数。
#分布式训练 #梯度归一化 #PyTorch DDP
讨论