深度学习加速必备知识原理级讲解：混合精度训练

最新推荐文章于 2025-12-05 11:17:30 发布

原创最新推荐文章于 2025-12-05 11:17:30 发布 · 1k 阅读

26 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #算法 #深度学习 #机器学习 #transformer #神经网络 #语言模型

第一部分：数字的物理形态（The Physics of Numbers）

在讨论训练之前，我们必须先钻到计算机的底层：计算机是如何存储小数的？

深度学习本质上就是大量的矩阵乘法和加法。如果你能减少每个数字占用的比特（Bit）数，你就能：

减少显存占用：存下更大的模型或 Batch Size。
提高带宽利用率：单位时间内传输更多数据。
加快计算速度：低精度计算单元（如 Tensor Core）算得更快。

1. 地基概念：IEEE 754 浮点数标准

在计算机中，我们不用“定点数”存小数，而是用科学计数法的二进制形式，这被称为浮点数（Floating Point）。

任何一个浮点数 V 在计算机里都由三部分比特位组成：

$V = (-1)^{S} \times 2^{E - \text{Bias}} \times (1 + M)$

这三个部分决定了数字的命运：

Sign (S, 符号位)：1位。决定是正数还是负数。
Exponent (E, 指数位)：决定数字的动态范围（Range）。也就是这个数最大能多大，最小能多小。指数位越多，能表示的数越大/越小。
Mantissa / Fraction (M, 尾数位/精度位)：决定数字的精度（Precision）。也就是小数点后能精确到多少位。尾数位越多，数字越精准，如果是钱，算得越细；如果是模型，梯度越准确。

2. 三位主角：FP32, FP16, BF16

混合精度训练，就是在这三种格式之间“反复横跳”。

(1) FP32 (Single Precision, 单精度) —— “全能选手”

这是深度学习默认的格式，也是最稳健的格式。

总长度：32 bits (4 Bytes)
结构：1位符号 + 8位指数 + 23位尾数
特点：
- 范围大：约 $\pm 3.4 \times 10^{38}$ 。
- 精度高：约 7个有效十进制数字。
- 缺点：太占内存，计算慢。

(2) FP16 (Half Precision, 半精度) —— “速度快但脆弱”

这是混合精度早期的主要优化对象。

总长度：16 bits (2 Bytes) —— 只有 FP32 的一半。
结构：1位符号 + 5位指数 + 10位尾数
致命弱点：
- 指数位太少（只有5位）：导致它的动态范围非常窄。
- 最大值只有 65504。超过这个数就会 Overflow（上溢，变成 Infinity）。
- 最小值（非正规数除外）约为 $6 \times 10^{-5}$ 。比这个更小的数（比如梯度中的 $10^{-7}$ ）会直接变成 0，这叫 Underflow（下溢）。
- 注：Underflow 是混合精度训练最大的敌人，后续我们讲 Loss Scaling 就是为了救它。

(3) BF16 (Brain Float 16) —— “为 AI 而生的变种”

这是 Google Brain 发明的，现在已成为 A100/H100 等新显卡的标配。

总长度：16 bits (2 Bytes)。
结构：1位符号 + 8位指数 + 7位尾数。
优点：
- 它直接把 FP32 的尾数砍掉了 16 位，但保留了和 FP32 一样的 8位指数。
- 优点：它的动态范围和 FP32 一模一样！不容易发生上溢或下溢。
- 缺点：精度低（只有7位尾数），但这对于深度学习（主要是统计学规律）通常是可以接受的。

3. 核心图解：为什么我们需要混合？

为了让你直观理解，我们做一个对比表：

格式	总位数	指数位 (决定范围)	尾数位 (决定精度)	动态范围 (约数)	核心问题
FP32	32	8	23	$10^{-38} \sim 10^{38}$	慢，显存占用大
FP16	16	5 (少!)	10	$6 \times 10^{-5} \sim 65504$	范围太窄，容易下溢变0
BF16	16	8 (同FP32)	7 (少!)	$10^{-38} \sim 10^{38}$	精度低，老显卡不支持

深度学习模型对精度（尾数位）其实不敏感。就像你识别一张猫的图片，像素点稍微有一点点噪点（精度损失），完全不影响你认出它是猫。

但是，模型对范围（指数位）非常敏感。特别是在计算梯度时，梯度通常非常小（例如 $0.000001$ ）。如果用 FP16，这个数字可能直接被当成 0 扔掉了，模型就学不到任何东西了。

这就引出了混合精度训练的核心矛盾：

我们想利用 16-bit 的速度（计算快、显存小），但又害怕它的范围窄导致梯度消失（Underflow）。

解决之道就是我们在第二部分要讲的：如何巧妙地结合 FP32 和 FP16，取长补短。

第二部分：混合精度的工作流 (The Workflow)

1. 支柱一：Master Weights (FP32 主权重)

这是混合精度训练中解决“精度不足”问题的关键。

问题场景： 假设我们在训练中进行权重更新（Weight Update）。公式很简单：

$W_{new} = W_{old} - \eta \cdot \nabla W$

其中 $\eta$ 是学习率（比如 $0.001$ ）， $\nabla W$ 是梯度。

在训练后期，梯度 $\nabla W$ 通常非常小（比如 $10^{-4}$ ），乘以学习率后，更新量可能只有 $10^{-7}$ 级别。

如果 $W$ 只是 FP16： FP16 的精度有限（尾数位少）。在计算机浮点数加法中，有一个经典的现象叫 "Swallowing"（大数吃小数）。如果一个大数（比如权重 $1.0$ ）加上一个极小的数（比如更新量 $0.0000001$ ），在 FP16 的精度下，结果可能依然是 $1.0$ 。

$1.0 \text{ (FP16)} + 0.0000001 \approx 1.0 \text{ (FP16)}$

这意味着：你算半天算出来的梯度，因为精度不够，根本加不进去！模型参数原地踏步。

解决方案 —— 维护一份 FP32 的“备份”：

我们始终在内存中保留一份完整的 FP32 权重（Master Weights）。
在做前向传播（Forward）和反向传播（Backward）时，我们把 FP32 权重转换成 FP16。
计算出的梯度也是 FP16 的。
关键一步：在更新权重时，我们将 FP16 的梯度转换回 FP32，然后更新到那个 FP32 的 Master Weights 上。

比喻：这就好比你要给一幅精细的油画（FP32）做修改。你先拍一张照片（转成 FP16），在照片上快速打草稿、涂改（计算梯度）。确认好怎么改之后，你拿起极细的画笔，回到原版油画（FP32）上进行微调。这样既快又不会丢失细节。

2. 支柱二：Loss Scaling (损失缩放)

这是混合精度训练中解决“范围不足（Underflow）”问题的关键。

问题场景： 我们在第一部分说过，FP16 最小只能表示约 $6 \times 10^{-5}$ 。然而，在深度学习（特别是 NLP 和 Transformer）中，梯度往往非常小，呈正态分布，很多都在 $10^{-7}$ 甚至更小。这些梯度一旦在反向传播中生成，如果用 FP16 存，直接变成 0。这叫 Gradient Underflow。

解决方案 —— 整体平移（Scaling）： 既然梯度太小，那我们在它变小之前，把它放大不就行了？

Scale Up：在前向传播结束，计算出 Loss（损失值）后，立刻把 Loss 乘以一个巨大的常数 S（Scaling Factor，比如 65536）。

$Loss' = Loss \times S$

Backward：根据链式法则，如果 $y = f(x)$ 变成 $S \cdot y$ ，那么梯度也会同步放大 $S$ 倍。

$\nabla W' = \nabla W \times S$

这使得原本极其微小的梯度（比如 $10^{-7}$ ），变成了较大的数值（ $10^{-7} \times 10^4 = 10^{-3}$ ），这就安全地落在了 FP16 的表示范围内（ $6 \times 10^{-5}$ 以上）。

Unscale：在更新权重之前，记得把这个 S 除回去，还原成真实的梯度值，然后再去更新 FP32 Master Weights。

可视化理解：想象梯度的数值分布在数轴的左侧（极小值区），FP16 的“可视窗口”在右侧。Loss Scaling 就是把整个梯度分布向右平移，强行拖进 FP16 的窗口里进行计算，算完再移回去。

3. 完整的工作流 (The Full Cycle)

把上面两个支柱结合起来，就是现代混合精度训练（AMP - Automatic Mixed Precision）的标准流程。请仔细看这个顺序，每一个环节都对应底层的逻辑：

准备阶段：

初始化模型，权重 W 存储为 FP32 (Master Weights)。

训练循环（Step）：

Weight Cast (FP32 -> FP16):

把 FP32 的权重复制一份，转换为 FP16。
- 目的：为了接下来的计算加速。
Forward Pass (FP16):

用 FP16 权重和 FP16 输入数据进行前向计算，得到预测值和 Loss。
- 优势：Tensor Core 全力开火，速度极快。
Loss Scaling (FP32):

计算出的 Loss (通常是 FP32) 乘以缩放因子 $S$ 。

Logits FP16 -> FP32 (为了防止 Softmax 溢出)。

Labels (标签) 也是 FP32 或 Int64。

计算 Loss: 在 FP32 下计算出 loss_value。

Loss Scaling: 在 FP32 下计算 scaled_loss = loss_value * S

$Loss_{scaled} = Loss \times S$
Backward Pass (FP16):

利用 $Loss_{scaled}$ 进行反向传播，计算梯度。
结果： 此时得到的梯度是 $\nabla W_{scaled}$ (FP16格式)，因为放大了，所以不会 Underflow。
Unscale Gradients (FP16 -> FP32):

将梯度 $\nabla W_{scaled}$ 转换回 FP32，并除以 $S$ 。

$\nabla W = \nabla W_{scaled} / S$
Weight Update (FP32):

利用真实的梯度 $\nabla W$ ，更新 FP32 Master Weights。

$W_{master} = W_{master} - \eta \cdot \nabla W$
（可选）更新 Scaling Factor S:

如果在这一步发现即使缩放了还是有梯度溢出（Inf/NaN），通常会减小 S跳过这一步；如果长期没溢出，会尝试增大 S。这叫动态损失缩放 (Dynamic Loss Scaling)。

在 PyTorch 的 AMP (Automatic Mixed Precision) 中，数据流是这样的：

Loss 计算 (FP32)：

模型输出（Logits）可能原本是 FP16，但在进入 Loss 函数（如 CrossEntropy）前会被转为 FP32，算出 loss (FP32)。
Scaling 乘法 (FP32)：

执行 scaled_loss = loss * scale_factor。这里输入、输出、中间计算全都是 FP32。
开始反向传播 (进入 FP16)：

调用 scaled_loss.backward()。

虽然起点是 FP32，但随着梯度流回网络内部（Backpropagation），为了加速，计算图中的中间梯度计算（矩阵乘法）会利用 Tensor Core 进行 FP16 运算。
梯度 Unscale (FP32)：

反向传播结束拿到梯度后，要把梯度除以 S还原。这个除法操作也是在 FP32 下进行的（通常是将梯度转回 FP32 后再除），以保证更新到 Master Weights 时的精度。

第三部分：代码实战与硬件加速 (Implementation & Hardware)

1. 硬件加速：Tensor Core 的秘密

你可能会疑惑：“把 FP32 转成 FP16，不就是少存点数据吗？为什么速度能快几倍？”

答案在于 NVIDIA 显卡里一种特殊的计算单元：Tensor Core。

普通 CUDA Core vs. Tensor Core

CUDA Core (FP32): 传统的计算单元。一次算一个数的乘法 ( $a \times b$ )。
Tensor Core (Mixed Precision): 它是为矩阵运算而生的暴力猛兽。它不像会计，像个流水线工厂。它在一个时钟周期内，不是算 1 个乘法，而是直接算一个 4x4 矩阵的乘加运算。

这里的“混合”发生在硬件内部

Tensor Core 的工作方式完美对应了我们讲的混合精度原理：

$D = A \times B + C$

输入 (A, B): 必须是 FP16。这意味着数据从显存搬运到芯片的速度快了一倍（带宽优势）。

运算与累加 (Accumulation): 虽然输入是 16位，但在 Tensor Core 内部，做矩阵乘法的“乘加”操作时，它会把结果累加到 FP32 的寄存器里！

输出 (D): 最终结果可以是 FP16 或 FP32。

底层视角：即使你告诉 PyTorch 用 FP16 训练，在最密集的矩阵乘法（Matrix Multiplication）内部，NVIDIA 的硬件其实悄悄帮你用 FP32 做了累加，防止了中间结果的精度损失。这就是为什么混合精度既快（输入吞吐大）又准（内部累加精度高）。

2. PyTorch 代码实战 (AMP)

以前我们要自己手写 FP32 转换和 Loss Scaling，很痛苦。现在 PyTorch 提供了 AMP (Automatic Mixed Precision) 工具箱。

虽然代码只有几行，但每一行都对应着我们之前讲的原理。

假设我们要训练一个简单的网络：

import torch
import torch.nn as nn
import torch.optim as optim

# 0. 准备工作
model = MyModel().cuda()
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# === 核心组件 1: GradScaler ===
# 它的作用就是管理 Loss Scaling (我们在第二部分讲的那个 S)
# 它会自动检测梯度是否溢出，并动态调整 S 的大小
scaler = torch.cuda.amp.GradScaler()

for input, target in data_loader:
    optimizer.zero_grad()
    
    # === 核心组件 2: autocast 上下文 ===
    # 它的作用是自动把原本 FP32 的运算（如卷积、矩阵乘）转成 FP16
    # 只有在这个范围内的操作才会走 Tensor Core 加速
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        output = model(input)           # 前向传播 (FP16)
        loss = criterion(output, target) # Loss计算 (内部自动转FP32)

    # === 核心组件 3: 缩放 Loss 并反向传播 ===
    # 回忆第二部分：loss.backward() 前要先乘以 S
    # 这里 scaler.scale(loss) 就是执行 loss * S (FP32运算)
    scaler.scale(loss).backward()

    # === 核心组件 4: Unscale + Update ===
    # scaler.step(optimizer) 做两件事：
    # 1. Unscale: 把梯度除以 S 还原
    # 2. Check Inf: 检查梯度有没有无穷大(Inf)或NaN
    #    - 如果有 Inf: 忽略这次更新 (optimizer.step() 不会被调用)
    #    - 如果没 Inf: 调用 optimizer.step() 更新 Master Weights
    scaler.step(optimizer)

    # === 核心组件 5: 更新 Scaling Factor ===
    # 根据这一步有没有发生溢出，决定下一步 S 是变大还是变小
    scaler.update()

3. 进阶：BF16 —— 抛弃 Scaler 的未来

既然是做大模型（LLM）研究的，必须知道 BF16 (Brain Float 16)。

我们在第一部分说过，BF16 的指数位和 FP32 一样多（8位），所以它不容易发生 Underflow/Overflow。

这对代码意味着什么？ 意味着如果我们用 BF16，我们就不再需要 Loss Scaling 了！

因为梯度不会变成 0，也不会溢出，所以 GradScaler 是多余的。这让训练更加稳定，不像 FP16 那样如果不小心 Scaler 没调好就容易崩。

BF16 代码实现 (需要 A100/H100/3090/4090 等较新显卡):

# 只需要改一行：dtype=torch.bfloat16
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    output = model(input)
    loss = criterion(output, target)

# 不需要 scaler.scale
loss.backward()

# 不需要 scaler.step
optimizer.step()

这也是为什么现在的 LLM（如 LLaMA, GPT-4）训练基本都全面转向 BF16 的原因。它既有 FP16 的速度，又有 FP32 的稳健性，省去了 Loss Scaling 的麻烦。

混合精度训练 · 全文总结

为了彻底理解这个概念，请脑海中保留这三张图：

物理层 (Data):
- FP32: 精准但臃肿的胖子。
- FP16: 只有 5 位指数的瘦子，容易“营养不良”（Underflow 变 0）。
- BF16: 砍断了腿（精度）但保留了上半身（范围）的变异体，专为 AI 优化。
逻辑层 (Workflow):
- Master Weights: 这种“主副本”机制，保证了微小的梯度更新不会因为精度不够被吞没。
- Loss Scaling: 这种“放大再缩小”的机制，像显微镜一样，把 FP16 看不见的微小梯度强行拉入可视范围。
硬件层 (Hardware):
- Tensor Core: 显卡里的特种部队。它吃进去的是粗糙的粮食 (FP16)，但在肚子里消化时用的是精细的工艺 (FP32 Accumulation)，从而实现了速度与精度的完美平衡。