PyTorch混合精度训练实战：如何用AMP让你的旧显卡也能跑大模型

PyTorch混合精度训练实战：如何用AMP让你的旧显卡也能跑大模型

手里只有一张RTX 3060，甚至更老的RTX 2060，看着动辄数十GB显存需求的大模型，是不是觉得训练任务遥不可及？别急着升级硬件，也别轻易放弃。很多时候，限制我们的不是显卡的绝对性能，而是对显存利用效率的认知。混合精度训练，正是打开这扇门的钥匙。它不是什么遥不可及的实验室技术，而是PyTorch框架内置的、经过实战检验的成熟方案。今天，我们不谈空洞的理论，直接上手，看看如何用PyTorch原生的AMP（Automatic Mixed Precision）工具，让你手头的消费级显卡也能承担起大模型的训练任务。

我见过太多开发者，一遇到显存不足就本能地调小batch size，或者干脆放弃更复杂的模型结构。这其实是一种资源的浪费。现代GPU，即便是消费级产品，其Tensor Core单元对半精度浮点数（FP16）的计算吞吐量，往往是单精度（FP32）的2到8倍。这意味着，合理利用FP16，不仅能省下近一半的显存，还能大幅提升计算速度。关键在于，如何安全、稳定地使用它，避免因精度损失导致的模型不收敛或数值溢出。这就是AMP要解决的问题。

1. 混合精度训练的核心原理：不只是省显存

在深入代码之前，我们有必要花几分钟理解混合精度训练到底在做什么。这能帮助你在遇到问题时，知道该从哪里排查。

深度学习训练中的数值主要有三种：模型权重（Weights）、激活值（Activations） 和梯度（Gradients）。传统训练全程使用FP32，每个数占用4字节。FP16则只占2字节，直观上能节省一半的存储空间。但直接全部换成FP16行不通，主要因为两个问题：

1. 数值范围不足：FP16能表示的最大正值约为65504，最小正值约为5.96e-8。而深度学习中的梯度值通常非常小，可能在1e-7甚至更小。在反向传播的链式乘法中，这些微小梯度用FP16表示时会直接下溢（Underflow）为零，导致权重无法更新。
2. 舍入误差：当权重值很大（例如1.0），而梯度更新量很小（例如1e-5）时，FP16可能无法精确表示这个微小的变化，导致更新无效。

AMP的聪明之处在于“混合”。它并非简单地将所有计算转为FP16，而是采用了一个主权重（Master Weights） 的概念。具体流程如下：

• 前向传播（FP16）：使用FP16的权重副本和输入数据进行计算，得到FP16的损失。
• 损失缩放（Loss Scaling）：将FP16的损失乘以一个缩放因子（如1024或2048），将其“放大”到FP16能较好表示的数值范围内。这一步是关键，它保护了那些微小的梯度不被淹没。
• 反向传播（FP16）：计算得到放大后的FP16梯度。
• 梯度反缩放与更新（转FP32）：将放大后的FP16梯度转换回FP32，并除以相同的缩放因子，得到真实的FP32梯度。然后用这个FP32梯度去更新保存在FP32精度下的主权重。
• 同步权重：将更新后的FP32主权重转换为FP16副本，用于下一次前向传播。

这个过程就像一个精密的放大器：在前向和反向计算这些“重活”上用FP16来提速省显存，而在关键的权重更新这个“精细活”上，用FP32来保证数值稳定和精度。

注意：并非所有计算都适合用FP16。像Softmax、LayerNorm这类对数值范围敏感的操作，强制使用FP16容易导致溢出或精度灾难。AMP内部维护了一个“白名单”和“黑名单”，自动为不同类型的操作分配合适的精度。

为了更直观地理解AMP模式下显存的分布变化，我们来看一个简单的对比表格：