神经网络Transformer架构中的混合精度训练 基础概念：计算精度与数值格式 在计算机中，数字以有限精度的二进制格式存储。深度学习训练中，最常用的是单精度浮点数（FP32），它使用32位比特表示一个数字，具有约7位有效十进制数字，能表示极大和极小的数值范围。混合精度训练的核心是引入半精度浮点数（FP16），它仅使用16位比特，范围更小，精度更低（约3-4位有效数字），但占用内存减半，计算速度（在支持硬件上）可大幅提升。 动机与挑战：为何及如何混合精度 直接使用FP16进行整个训练过程会带来两个主要问题。一是 数值下溢 ：梯度值通常非常小，可能小于FP16能表示的最小正值，导致变为零，使训练停滞。二是 数值上溢 ：在大型模型（如Transformer）中，前向传播的激活值或反向传播的梯度值可能过大，超出FP16的表示范围，变成无穷大（NaN），导致训练崩溃。混合精度训练通过精细的管理策略来克服这些问题，以在保持训练稳定性和最终模型精度的前提下，利用FP16的速度和内存优势。 核心技术组件：损失缩放 这是解决梯度下溢问题的关键。算法会自动监测梯度的幅度。如果发现梯度值普遍过小，它会将一个缩放因子（例如，256、512、1024等）乘到前向传播计算出的损失值上。根据链式法则，这个放大的损失值在反向传播时会使所有梯度按相同比例放大，从而将它们“抬升”到FP16能够有效表示的安全范围内。在优化器更新权重之前，这些放大后的梯度会被正确地缩小回原来的比例。 核心技术组件：主权重副本与精度转换 这是确保模型最终精度与FP32训练相当的核心。在混合精度训练中，模型权重在内存中维护两个副本： FP16副本 ：用于前向传播和反向传播计算，享受FP16的计算速度优势。 FP32主副本 ：用于优化器更新步骤。在每次反向传播计算出FP16梯度后，梯度被转换成FP32，用于更新这个FP32主权重。在下一次迭代前，将更新后的FP32主权重再转换为FP16，用于下一次前向传播。这样做是因为优化器更新（如动量、权重衰减）需要高精度来累积微小的更新量，避免舍入误差累积导致训练漂移。 工作流程与硬件协同 一个完整的混合精度训练迭代步骤如下： a. 前向传播 ：使用FP16权重和输入数据进行计算，得到FP16的激活和损失。 b. 损失缩放 ：将计算出的损失乘以一个动态或静态的缩放因子。 c. 反向传播 ：使用FP16权重和缩放后的损失，计算得到FP16梯度。 d. 梯度缩放 ：将FP16梯度除以相同的缩放因子，将其恢复为真实值（但此时梯度已处于FP16安全范围内）。 e. 优化器更新 ：将FP16梯度转换为FP32，并用其更新存储在内存中的FP32主权重副本。 f. 权重同步 ：将更新后的FP32主权重转换为FP16，覆盖原有的FP16权重副本，为下一次迭代做好准备。 现代GPU（如NVIDIA的Tensor Core）和AI加速器专门优化了FP16计算，并能高效处理这种FP16计算与FP32累加的组合操作。 优势、应用与实现 混合精度训练的主要优势是 显著减少GPU显存占用 （可训练更大模型或使用更大批次）和 提升训练吞吐量 （加快训练速度），通常能达到1.5倍至3倍的加速，而对最终模型精度的影响微乎其微。它已成为训练大型Transformer模型（如GPT、BERT等）及计算机视觉模型的 标准实践 。主流深度学习框架（如PyTorch、TensorFlow）都通过API（如 torch.cuda.amp ）提供了自动化的混合精度训练支持，开发者只需添加少量代码即可启用。