福桃九分饱FP16一开就loss炸了，TF32到底要不要开，大模型微调怎么把显存降下来，这类问题最近被反复问到。大模型训练和微调的成本越来越高，大家第一个想动手优化的就是混合精度训练，想用FP16配合TF32把显存占用和算力消耗压下来。这个思路是对的，尤其是当数据以token形式进入模型时，每一个token都要完整走一遍前向和反向，显存和算力消耗几乎跟token总数挂钩。这里我把自己踩过的坑、验证过的方案都摊开讲一讲，希望能帮你把这套流程跑通。
1. 混合精度训练到底在解决什么问题？
1.1 FP16、FP32、TF32的关系先搞明白
想用好混合精度训练，第一步不是写代码，是把浮点格式搞明白。一个浮点数在计算机里由符号位、指数位、尾数位组成，位数不同，能表达的范围和精度就不同：
- FP32：1位符号 + 8位指数 + 23位尾数，总共32位。这是深度学习之前默认的精度，也叫单精度。
- FP16：1位符号 + 5位指数 + 10位尾数，总共16位。指数位只有5位，导致它的动态范围特别小，最大只能到65504，最小正常数是6.1e-5左右。
- TF32：这个最容易被误会，很多人以为它是32位格式，实际上它是NVIDIA Ampere架构开始引入的Tensor Core计算格式，内部是1位符号 + 8位指数 + 10位尾数，总共19位，但在32位容器里存储和处理。指数范围跟FP32一致，所以动态范围很大，代价是尾数只有10位，做乘法时有效精度其实和FP16一个水平。
- BF16：1位符号 + 8位指数 + 7位尾数，总共16位。指数位和FP32完全一样，动态范围合适，但尾数位只有7位，有效精度比FP16还低。
很多训练卡在FP16上，就是因为低估了FP16的窄范围。打个比方：FP32像一把量程30厘米、刻度到毫米的尺；FP16量程缩到6厘米，刻度还是十分之一毫米；BF16量程够大但刻度变粗；TF32是量程大、刻度中等偏粗的一个变种。正常数值在FP32里没问题，切到FP16后要么溢出变成inf，要么小数部分直接归零，一下就把训练弄崩了。
1.2 为什么说省显存就能降低token消耗
模型训练里，输入文本被切分成token后，通常以[batch_size, seq_len]的索引矩阵进入网络，经过embedding后变成[batch_size, seq_len, hidden_size]的高维张量。从这一层开始，每一层注意力机制、每一层FFN的输入输出和中间激活，规模都跟batch内的token总数成正比。
也就是说，在显存受限的情况下，能塞进GPU的token总数决定了训练效率。想要多塞token，无非两条路：增大batch size，或者拉长序列长度，但都需要更多显存。混合精度训练把权重和中间激活的一部分从32位砍到16位，同等显存预算下能直接多装一批token。这里的"减少token消耗"，指的是模型每处理一个token所占用的显存和算力成本下降了，并不是说模型输出变短了。
这个账在推理场景更好算。多轮对话或者长上下文场景里，模型会为每个已经出现的token缓存一组Key和Value，也就是KV Cache。假设hidden_size是4096，模型有32层，用FP32存KV Cache的话，每个token每层要存两个向量，大约4096×2×2×32=512KB。同样的模型切成FP16，只算KV Cache一个token就能省256KB。上下文一旦拉到几万token，差距就非常明显了。
1.3 Tensor Core真正把算力红利吃进去
省显存只是混合精度的一半好处，另一半是速度。