PyTorch Lightning中混合精度训练与手动优化的常见陷阱

在PyTorch Lightning项目中使用混合精度训练(16位混合精度)时，开发者可能会遇到一个典型错误："Attempted unscale_ but _scale is None"。这个错误通常发生在尝试手动优化多个优化器的场景下。

问题现象

当开发者配置trainer.precision='16-mixed'并尝试手动管理多个优化器时，系统会抛出上述错误。具体表现为：

• 使用3个Adam优化器
• 每个优化器都配有ReduceLROnPlateau学习率调度器
• 在默认精度(32位)下运行正常
• 切换到16位混合精度时出现错误

根本原因

问题的核心在于混合精度训练需要特殊处理梯度缩放。PyTorch Lightning的自动混合精度(AMP)功能依赖于梯度缩放器(GradientScaler)来防止16位浮点数下的数值下溢。

在手动优化模式下，开发者错误地直接调用了loss.backward()而不是Lightning提供的manual_backward(loss)方法。后者会正确处理AMP所需的梯度缩放步骤。

解决方案

正确的做法是：

1. 在LightningModule中重写optimizer_step方法时
2. 使用self.manual_backward(loss)替代原生的loss.backward()
3. 确保所有梯度计算都通过Lightning提供的接口

这种方法确保了：

• 梯度缩放器被正确初始化
• 前向传播和反向传播的精度转换得到妥善处理
• 多个优化器之间的梯度更新协调一致

最佳实践

对于需要在PyTorch Lightning中手动管理多个优化器的场景，建议：

1. 始终使用Lightning提供的抽象接口(如manual_backward)
2. 在混合精度训练时特别注意梯度缩放
3. 测试时先使用默认精度验证逻辑正确性
4. 切换到混合精度后监控数值稳定性
5. 考虑使用Lightning的自动优化功能(如果可以满足需求)

理解这些底层机制可以帮助开发者更有效地利用PyTorch Lightning的高级功能，同时避免常见的陷阱。