第4章混合精度训练

第四章：混合精度训练——FP16与FP32的实战艺术

各位同学，今天我们来聊聊混合精度训练。说实话，这玩意儿刚出来的时候，我是持怀疑态度的。FP16？那精度够用吗？直到我在一个高频交易模型上，把训练时间从三天压缩到十个小时，我才真正服了。

混合精度训练，说白了就是让模型在训练时同时使用FP16和FP32两种精度。FP16算得快、省显存，FP32保精度、稳如狗。两者配合，效果出奇的好。

FP16与FP32：精度与速度的博弈

先看一组数据，这是我实测的对比：

精度类型	存储大小	计算速度	数值范围	精度
FP32	4字节	1x	~3.4e38	高
FP16	2字节	2-8x	~6.5e4	中
BF16	2字节	2-8x	~3.4e38	中

看到没？FP16的数值范围只有6.5万，而FP32是3.4亿。这意味着什么？意味着梯度很容易溢出或者下溢。我在项目中遇到过，一个简单的LSTM模型，用纯FP16训练，loss直接变成NaN。嗯，这就是典型的梯度爆炸。

核心要点：FP16的指数位只有5位，而FP32有8位。这3位的差距，就是精度和范围的巨大鸿沟。

AMP自动混合精度：让框架替你操心

AMP（Automatic Mixed Precision）是NVIDIA推出的解决方案。它的核心逻辑很简单：

权重和梯度用FP16存储
主副本用FP32保存
损失计算和敏感操作自动切换到FP32
动态调整损失缩放因子

我个人习惯用PyTorch的AMP实现，代码非常简洁：

import torch
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

你想想看，就这几行代码，训练速度能提升2-3倍。我第一次用的时候，还特意对比了纯FP32的结果，精度几乎没差别。

损失缩放：防止梯度消失的杀手锏

为什么需要损失缩放？因为FP16能表示的最小正数是6e-8。如果梯度小于这个值，直接变成0。这就是梯度消失的另一种形式。

损失缩放的做法是：在反向传播前，把损失值乘以一个缩放因子（比如2^16），让梯度落在FP16的可表示范围内。反向传播完成后，再把梯度除以这个因子。

实战技巧：我一般把初始缩放因子设为2^16，然后让AMP动态调整。如果连续N步没有溢出，就增大因子；如果出现溢出，就减小。这个N我通常设为2000。

我曾经在一个Transformer模型上，因为忘记做损失缩放，训练了整整两天才发现梯度全是0。嗯，那种感觉，就像你跑了马拉松却发现终点线画错了地方。

梯度裁剪：防止梯度爆炸的保险丝

梯度裁剪和混合精度训练是绝配。为什么？因为FP16的数值范围小，一旦梯度爆炸，直接溢出变成NaN。而梯度裁剪就像保险丝，在电流过大时自动断开。

我常用的梯度裁剪方式有两种：

按值裁剪：把梯度限制在[-max_value, max_value]之间
按范数裁剪：把梯度的L2范数限制在max_norm以内

代码实现也很简单：

# 按范数裁剪（推荐）
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 按值裁剪
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

注意：梯度裁剪要在损失缩放之后、optimizer.step()之前执行。顺序错了，效果天差地别。

实战案例：量化交易LSTM模型优化

好了，理论说完了，咱们来点实际的。这是我优化一个5分钟K线预测模型的过程：

原始配置：

模型：3层LSTM + 2层全连接
参数：约5000万
训练数据：5年美股分钟级数据
硬件：单卡RTX 3090

优化前：纯FP32训练，一个epoch需要4小时，显存占用22GB。

优化后：混合精度训练，一个epoch只需要1.5小时，显存占用降到12GB。

具体配置如下：

import torch
from torch.cuda.amp import autocast, GradScaler

# 初始化
scaler = GradScaler(init_scale=2**16)
model = LSTMModel(input_dim=50, hidden_dim=256, num_layers=3).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)

# 训练循环
for epoch in range(100):
    for batch in dataloader:
        optimizer.zero_grad()
        
        with autocast():
            pred = model(batch['features'])
            loss = criterion(pred, batch['labels'])
        
        scaler.scale(loss).backward()
        
        # 梯度裁剪
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        scaler.step(optimizer)
        scaler.update()

这里有个细节：scaler.unscale_()这步很多人会忘。它的作用是把缩放后的梯度还原，这样梯度裁剪才能正确工作。我曾经因为这个bug，排查了整整一个下午。

知识体系总览

下面这张图，是我对混合精度训练核心逻辑的总结：

这张图把整个流程串起来了。你想想看，从FP32主权重开始，到FP16前向传播，再到损失缩放、反向传播，最后梯度裁剪和更新。每一步都有它的道理。

我的建议：刚开始用混合精度训练时，先在小模型上验证。跑几个epoch，对比FP32的结果。确认精度没问题后，再上大模型。我曾经因为太自信，直接在大模型上开AMP，结果训了三天发现loss曲线不对。嗯，从那以后，我学会了先做小规模验证。

好了，混合精度训练的核心内容就这些。记住：FP16负责速度，FP32负责精度，AMP负责协调，损失缩放和梯度裁剪负责兜底。四者配合，才能发挥最大威力。

📚 课程目录

第四章：混合精度训练——FP16与FP32的实战艺术

FP16与FP32：精度与速度的博弈

AMP自动混合精度：让框架替你操心

损失缩放：防止梯度消失的杀手锏

梯度裁剪：防止梯度爆炸的保险丝

实战案例：量化交易LSTM模型优化

知识体系总览