第四章：数据类型优化——float16、混合精度与bfloat16

说到显存优化，数据类型这块是性价比最高的切入点。我刚开始做大模型训练时，显存动不动就爆，后来发现很多地方根本不需要float32那么高的精度。说白了，你用高精度存了一堆小数点后十几位的数字，但模型根本用不上那么多信息。

这一章，我带你看看float16、混合精度训练（AMP）和bfloat16到底怎么用。嗯，都是实战中反复踩坑总结出来的经验。

4.1 为什么float16能省一半显存？

先看个简单的对比：

数据类型	占用字节	指数位	尾数位	表示范围
float32	4字节	8位	23位	~3.4e38
float16	2字节	5位	10位	~65504
bfloat16	2字节	8位	7位	~3.4e38

float16直接砍掉一半字节。你想想看，一个7B参数的模型，光参数就占28GB（float32）。换成float16，直接降到14GB。这省出来的空间，够你塞下更大的batch size了。

但有个坑——float16的表示范围太小了。最大值才65504，而float32能到3.4e38。我在项目中遇到过梯度爆炸，float16直接溢出变成NaN，整个训练就崩了。所以不能无脑换。

4.2 混合精度训练（AMP）——我常用的方案

混合精度训练，说白了就是「该用float16的地方用float16，该用float32的地方保留float32」。PyTorch的AMP（Automatic Mixed Precision）帮我省了不少事。

核心思路就三点：

前向传播：用float16计算，省显存、加速
反向传播：梯度也用float16，但关键地方做缩放
权重更新：主权重保留float32副本，防止精度丢失

代码示例：

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

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()  # 梯度缩放器

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():  # 自动混合精度上下文
        output = model(data)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()  # 缩放梯度
    scaler.step(optimizer)         # 反缩放并更新
    scaler.update()                # 更新缩放因子

这里有个细节：GradScaler会自动调整缩放因子，防止梯度下溢。我刚开始用的时候没加这个，结果小梯度全被截断成0了，模型根本学不动。

关键点：AMP不是简单地把所有东西都转成float16。它会在计算过程中动态判断哪些操作需要高精度，哪些可以用低精度。比如LayerNorm和Softmax这些操作，AMP会自动用float32计算。

4.3 bfloat16——Google的杀手锏

bfloat16是Google搞出来的，后来被NVIDIA的Ampere架构支持。它和float16一样占2字节，但指数位有8位，和float32一样多。

这意味着什么？

表示范围：和float32一样大，不会溢出
精度：尾数只有7位，比float16的10位少
适用场景：大模型训练，尤其是梯度容易爆炸的场景

我个人习惯在训练LLM时优先用bfloat16。为什么呢？因为大模型的梯度范数经常很大，float16动不动就溢出，而bfloat16完全扛得住。

代码示例：

# 如果你的GPU支持bfloat16（A100、H100等）
model = model.bfloat16()

# 或者用AMP指定bfloat16
with autocast(dtype=torch.bfloat16):
    output = model(input)

小技巧：bfloat16不需要GradScaler。因为它范围够大，不会下溢也不会溢出。代码更简洁，训练更稳定。

4.4 避坑指南——我踩过的雷

我曾经在训练一个BERT-large模型时，直接全部用float16，结果loss死活不降。查了半天，发现是Embedding层的梯度太小，被float16截断了。

后来我总结了几条经验：

不要全盘float16：BatchNorm、LayerNorm这些操作，保留float32
监控梯度范数：如果发现梯度突然变成0或NaN，八成是精度问题
bfloat16不是万能药：虽然范围大，但精度低。对精度敏感的任务（比如科学计算），还是得用float32
检查GPU支持：老显卡（V100之前）不支持bfloat16，别白费功夫

警告：混合精度训练时，如果发现loss曲线异常抖动，先检查一下缩放因子。我曾经遇到过scaler.update()没写对，缩放因子一直不更新，模型训练了三天全白费。

4.5 知识体系图

下面这张图帮你理清数据类型优化的核心逻辑：

嗯，这张图把三种方案的优缺点和适用场景都串起来了。你想想看，实际项目中，我一般这样选：

模型小于1B参数：直接float16 + AMP，简单高效
模型大于1B参数：优先bfloat16，省心
GPU不支持bfloat16：AMP + float16，记得加GradScaler

最后说一句，数据类型优化是显存优化的第一步，也是最容易见效的一步。别小看这2字节的差距，在7B模型上就是14GB的显存差距。我见过太多人一上来就搞模型并行、流水线并行，其实先把数据类型调好，很多问题就解决了。

想要更多资料，联系可可

📚 课程目录

第四章：数据类型优化——float16、混合精度与bfloat16

4.1 为什么float16能省一半显存？

4.2 混合精度训练（AMP）——我常用的方案

4.3 bfloat16——Google的杀手锏

4.4 避坑指南——我踩过的雷

4.5 知识体系图