混合精度训练是一种深度学习训练技术，它结合使用不同数值格式（主要是 FP16 或 BF16 与 FP32）来执行模型的不同部分计算。其目标是在保持模型收敛速度和精度的同时，显著提高训练速度并减少显存占用。

1. 浮点精度回顾：FP32 vs. FP16

1.1 FP32（单精度浮点数）

• 全称： Single-Precision Floating-Point (32-bit)

• 结构： 1 位符号位 (Sign) + 8 位指数位 (Exponent) + 23 位尾数位 (Mantissa)。

• 特点： 具有较大的动态范围和较高的精度。在深度学习普及之前，所有训练都使用 FP32 以确保数值稳定。

• 缺点： 占用 32 位（4 字节）存储空间，计算速度相对较慢。

1.2 FP16（半精度浮点数）

• 全称： Half-Precision Floating-Point (16-bit)

• 结构： 1 位符号位 + 5 位指数位 + 10 位尾数位。

• 特点：

  • 动态范围小： 只有 5 位指数位，能表示的数值范围比 FP32 小得多。

  • 精度低： 只有 10 位尾数位，精度比 FP32 低得多。

• 优点： 占用 16 位（2 字节）存储空间，数据传输量减半，并且计算速度极快（尤其是在 Tensor Core 上）。

2. 混合精度训练的工作原理

混合精度训练的核心思想是扬长避短：用 FP16 的速度优势来加速计算，同时用 FP32 的稳定优势来保护关键操作。

2.1 训练流程（NVIDIA APEX 或 PyTorch AMP）

1. 存储（FP32 Master Copy）： 模型的权重（Weights）和偏置（Biases）始终以 FP32 格式存储一份主拷贝（Master Copy）。

2. 前向传播（FP16/BF16）：

   • 在计算前，将 FP32 主拷贝权重转换为 FP16。

   • **大部分计算（如矩阵乘法、卷积）**都使用 FP16 进行，以利用 Tensor Core 的速度。

3. 损失值计算（FP32）： 损失值（Loss）和梯度（Gradients）的累加通常转回 FP32 进行，以确保精度。

4. 梯度缩放（Loss Scaling）：

   • 解决下溢问题： 由于 FP16 的动态范围小，梯度的绝对值通常很小，很容易四舍五入到零（下溢）。

   • 解决方案： 在计算损失后，将损失值放大一个固定的比例因子（如 $2^N$），使得梯度在反向传播时保持在 FP16 的可表示范围内。

5. 权重更新（FP32）：

   • 将缩放后的 FP16 梯度缩小（除以相同的比例因子）。

   • 使用缩小的梯度和 FP32 主拷贝进行权重更新。这保证了模型参数的长期稳定性。

3. Tensor Core：加速的核心引擎

Tensor Core 是混合精度训练成功的关键。

3.1 专门的 FP16/BF16 硬件

Tensor Core 是 GPU 上的专用硬件单元，它执行 $D=A\times B+C$ 操作时，要求 $A$ 和 $B$ 是 FP16/BF16 或其他低精度格式。

• 加速原理： Tensor Core 的设计使其在处理 FP16 数据时，比标准 CUDA 核心在处理 FP32 数据时，能实现指数级的吞吐量提升（例如，在 Ampere 架构上，TF32 的性能是 FP32 的 8 倍，FP16/BF16 的性能是 FP32 的 16 倍）。

3.2 对 TF32 的支持（Ampere 及后续架构）

自 Ampere 架构（A100）以来，NVIDIA 引入了 TensorFloat-32 (TF32) 格式：

• 特点： TF32 拥有 FP32 的动态范围，但精度接近 FP16。

• 训练简化： 如果你的代码仍然使用 FP32 类型，Tensor Core 可以自动在硬件层面将 FP32 数据转换为 TF32，然后使用 Tensor Core 进行加速。这使得开发者能够轻松地获得性能提升，而无需手动管理 FP16 转换和梯度缩放。

4. 混合精度训练的优势总结

结论：

混合精度训练已成为现代深度学习的默认最佳实践。它通过 Tensor Core 的硬件加速和精妙的数值管理策略（如梯度缩放），成功地结合了 FP16 的速度优势和 FP32 的数值稳定性，是推动大规模 AI 模型（如 Transformer）高效训练的关键技术之一。