混合精度训练调优实战 - 第3节

混合精度训练原理：Master Weights与权重更新

混合精度训练，说白了就是让模型用FP16算得快，但还得保证精度不掉。这里有个核心矛盾：FP16能表示的数字范围太小了。我刚开始接触时也纳闷，既然FP16这么容易溢出，为什么大家还要用？

答案就在Master Weights里。

什么是Master Weights？

Master Weights，就是一份FP32精度的权重副本。模型前向和反向用FP16算，但权重更新的时候，我们偷偷保留一份FP32的版本。

为什么要这么做？你想想看，FP16的梯度如果直接更新权重，很容易就变成0了。特别是那些小梯度，FP16根本表示不了。我在项目中遇到过，一个BERT-large模型训练到一半，loss突然不降了。查了半天，发现是梯度太小，被FP16截断成了0。

具体流程是这样的：

1. 维护一份FP32的Master Weights
2. 前向传播时，把Master Weights转成FP16
3. 用FP16做前向和反向计算
4. 拿到FP16的梯度后，转回FP32
5. 用FP32的梯度更新Master Weights

权重更新的精度管理

我习惯把权重更新比作「存钱」。FP16就像一个小钱包，只能装几百块。FP32就像银行卡，能存几百万。你用小钱包做交易（前向反向），但最终的钱还是存在银行卡里（Master Weights）。

代码实现其实不复杂：

# 伪代码示例
master_weights = model.parameters().float()  # FP32副本

for batch in dataloader:
    # 前向用FP16
    with torch.cuda.amp.autocast():
        loss = model(batch)
    
    # 反向传播
    scaler.scale(loss).backward()
    
    # 梯度转FP32，更新Master Weights
    scaler.step(optimizer)
    scaler.update()

嗯，这里要注意：optimizer里管理的参数必须是Master Weights，不是模型本身的FP16参数。我见过有人搞混，结果模型越训越差。

Loss Scaling（损失缩放）

Loss Scaling，这名字听着高大上，其实道理很简单。FP16能表示的最小正数是2^-24，大约6e-8。如果梯度比这个还小，直接就变0了。

怎么办？把loss放大啊。

我刚开始做混合精度时，总觉得这步是多余的。直到有一次训练一个图像分割模型，发现小目标的梯度全没了。嗯，从那以后我再也不敢跳过Loss Scaling了。

缩放策略

策略	说明	我推荐
固定缩放	手动设置一个缩放因子，比如1024	不推荐，容易过缩放或欠缩放
动态缩放	根据梯度是否溢出自动调整	强烈推荐，AMP默认就是这个
无缩放	直接硬跑FP16	别试，会后悔的

动态缩放是怎么工作的？

• 如果连续N步没有溢出，就把缩放因子翻倍（比如从1024到2048）
• 如果某一步溢出了，就回退到上一个安全的缩放因子
• PyTorch AMP默认N=2000，我个人觉得这个值挺合理的

自动混合精度（AMP）机制

AMP，全称Automatic Mixed Precision。说白了就是让框架帮你决定：哪些操作用FP16，哪些用FP32。

为什么不能全用FP16？因为有些操作对精度太敏感了。比如BatchNorm的均值和方差，用FP16算会出大问题。我踩过这个坑，训练一个ResNet-50，用了全FP16，结果验证集准确率掉了5个点。

AMP的精度管理策略

PyTorch的AMP主要做三件事：

1. 自动Cast：前向传播时，自动把输入转成FP16
2. 白名单机制：某些操作强制用FP32（比如Softmax、BatchNorm）
3. 黑名单机制：某些操作强制用FP16（比如卷积、矩阵乘法）

我画了一张图，帮你理解AMP的决策流程：

AMP的三种模式

PyTorch提供了三种精度模式，我整理了一下：

模式	精度	速度提升	适用场景
O0	纯FP32	1x	调试、精度验证
O1	混合精度（推荐）	1.5-2x	大部分训练任务
O2	几乎全FP16	2-3x	对精度不敏感的任务
O3	纯FP16	3x+	不推荐，除非你确定没问题

我个人习惯用O1模式。O2虽然快，但我在项目中遇到过几次精度下降的问题。O1是最稳妥的选择。

前向与反向传播中的精度管理

前向传播时，AMP会自动把输入转成FP16。但这里有个坑：如果你的模型里有自定义操作，AMP可能不知道该怎么处理。

我举个例子。有一次我写了一个自定义的损失函数，里面用了torch.log。AMP默认把这个操作也转成了FP16，结果log(0)直接变成了-inf。嗯，从那以后我学会了用@torch.cuda.amp.custom_fwd和@torch.cuda.amp.custom_bwd装饰器。

反向传播的精度管理

反向传播时，梯度是用FP16计算的。但梯度更新前，必须转回FP32。这个转换由GradScaler自动完成。

我曾经犯过一个错误：手动把梯度转成FP16再更新权重。结果模型直接不收敛了。后来才发现，GradScaler已经帮我做了这件事，我多此一举了。

正确的做法是：

scaler = torch.cuda.amp.GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    
    with torch.cuda.amp.autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    # 反向传播
    scaler.scale(loss).backward()
    
    # 梯度更新（自动处理精度转换）
    scaler.step(optimizer)
    scaler.update()

你看，代码其实很简单。但背后的原理值得你花时间理解。混合精度训练不是魔法，它是一套精心设计的精度管理策略。

混合精度训练的核心，就是在速度和精度之间找到平衡点。Master Weights保证权重更新的精度，Loss Scaling防止梯度消失，AMP自动管理操作精度。这三者配合，才能让FP16训练既快又稳。

想要更多资料，联系可可