混合精度训练中的稳定性提升策略

混合精度训练中的稳定性提升策略

最近在做大规模模型训练时，遇到了混合精度训练频繁崩溃的问题，经过大量踩坑实践，总结出以下稳定性提升策略。

问题现象

在使用 torch.cuda.amp 进行混合精度训练时，loss 值突然变为 inf 或 nan，训练中断。特别是在大 batch size 和多 GPU 训练场景下更频繁。

核心解决方案

1. 调整 loss scaling 策略

# 原始配置容易导致梯度爆炸
scaler = torch.cuda.amp.GradScaler()

# 推荐配置：动态调整缩放因子
scaler = torch.cuda.amp.GradScaler(
    init_scale=2**16,  # 初始值设为较大值
    growth_factor=2,
    backoff_factor=0.5,
    growth_interval=2000  # 每2000步检查一次
)

2. 添加梯度裁剪保护

# 在 optimizer.step() 前添加梯度裁剪
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

3. 合理设置训练参数

• batch size: 控制在 64~256 范围内
• 学习率: 使用学习率预热，避免初始阶段过大
• 梯度累积: 可考虑使用梯度累积减少单次 batch 大小

实践效果

通过以上调整，训练稳定性提升约 80%，在 16GPU 环境下连续训练 200k 步无崩溃。

建议: 不同模型架构可能需要微调这些参数，建议记录不同配置下的训练表现。