如果你自己从零训过超过十亿参数的大模型大概率经历过这些场景代码写好了数据也检查过了结果训练卡刚跑起来显存不足的报错迎面砸过来把模型换成低精度之后显存安静了loss 却开始抽风要么突然变成 nan要么几十个 step 过去纹丝不动。这两个问题根子往往都在同一个地方——你用什么精度的浮点数做训练。这篇文章想把 FP32、FP16 与混合精度训练这三件事彻底讲透。我会先拆解浮点数在内存里的真实结构说明为什么 FP16 能让训练变快、变省却又不能无脑替换 FP32然后讲清混合精度训练最常见的三种设计机制FP32 权重主副本、损失缩放、动态调整最后落到大模型训练的真实场景里算一算显存账并整理一份踩坑排查表。适合的人群是训练过模型但没细究过数值精度的工程师准备开始训练大模型的学生以及做推理加速但想理解模型数值表现的人。1. 为什么大模型必须正视浮点精度FP32 与 FP16 的底层差异1.1 一个浮点数在内存里的真实结构很多人把“FP16”简单理解成“比 FP32 省一半显存”但这个理解不够。要搞懂混合精度得先知道一个浮点数在内存里是怎么组织的。现代计算机的浮点格式基本都遵循同一套国际标准IEEE 754核心是用“符号位 指数位 尾数位”三部分表示一个数。你可以把它类比成科学计数法一个数 符号 × 尾数 × 2 的指数次方。符号位决定正负指数位决定数值的动态范围也就是能表示多大和多大的数尾数位决定精度也就是能精确到几位有效数字。两种最常见精度的结构差异如下类型符号位指数位尾数位最大有限值最小正常数FP321 bit8 bit23 bit约 3.4 × 10^38约 1.2 × 10^-38FP161 bit5 bit10 bit65504约 6.1 × 10^-5指数位 8 位 vs 5 位决定了 FP32 能表示从 10^-38 到 10^38 的夸张范围而 FP16 最大只能到 65504。这个数字在深度学习中是个隐患loss 或梯度一旦超过 65504FP16 里就会直接变成无穷大。尾数位 23 位 vs 10 位决定了有效数字的多少。23 位二进制大约对应 6.9 位十进制有效数字10 位二进制大约对应 3 位十进制有效数字。也就是说FP16 的精度只有 3 位左右一个数存进去后面那些小数位会被直接抹掉。举个具体例子假设某个权重当前数值在 1024 附近FP16 在这段区间的分辨能力只有 1 左右。等于说你要在这个权重上加一个 0.01 量级的梯度更新FP16 根本“看不见”这个变化更新直接无效。这种“大数吃小数”的舍入问题是 FP16 无法独立承担训练任务的第一个原因。1.2 FP16 的收益那么大代价藏在哪既然 FP16 精度这么低为什么大模型训练还非要用它因为收益实在太大了。第一是显存减半。一个 1000 亿参数的模型权重用 FP32 存储需要 400GB换成 FP16 只需要 200GB。大模型训练时显存是稀缺资源能省一半权重空间意味着可以塞下更大的 batch或者塞下更长的序列。第二是带宽减半。训练过程要反复读取权重和激活值数据搬运的时间在整个训练里占比很高。同样的数据量FP16 只需要搬一半字节通信开销也随之下降。第三是算力翻倍。当前主流加速卡都内置了针对半精度矩阵乘法的专用计算单元FP16 的矩阵运算吞吐通常是 FP32 的两倍甚至更多。大模型训练绝大部分时间花在矩阵乘法上这让 FP16 成为了实际上的“主力计算格式”。那你可能想问精度低这么多模型不会训崩吗这里有个关键观察深度学习训练本身对数值噪声有一定容忍度。梯度下降本来就是在噪声中逼近最优解小规模的舍入误差会被 batch 之间的随机波动部分抵消。真正让 FP16 无法直接用的不是“有点误差”而是“误差大到权重更新时被吃掉”以及“数值溢出到不可收拾”。1.3 上溢、下溢与舍入FP16 不能全信任的三个理由FP16 在实际训练中有三个典型问题必须单独拿出来讲。第一个是上溢Overflow。FP16 最大只能到 65504而反向传播里偶尔会出现很大的梯度值尤其在训练早期和 batch 特别“难”的时候。一旦超过 65504数值就变成 inf 或 nan之后整个训练都乱了。第二个是下溢Underflow。这个比上溢更隐蔽也更常见。FP16 能表示的带完整精度的最小正常数是 6.1 × 10^-5虽然下面还有更小的次正规数但精度已经非常差。而深度学习中大量梯度的数量级在 10^-6 到 10^-8 甚至更小这些梯度在 FP16 里直接变成 0。结果就是模型“看起来还在训练”但实际上参数已经不再更新了。第三个是舍入误差的累积。权重更新本质上是成千上万次小步的累积。如果每一步都在 FP16 里操作舍入误差会反复叠加更糟糕的是当权重数值较大而梯度更新量较小时更新会被直接吞掉。这就是前文说的“大数吃小数”。这三个问题单独靠“把模型切成 FP16”解决不了。所以业界才发展出了一套更精细的做法不是所有东西都用 FP16而是让 FP16 和 FP32 各司其职。这就是混合精度训练。2. 混合精度训练的三大设计支柱2.1 第一根支柱FP32 权重主副本让更新不“近视”混合精度训练最核心的思想用一句话概括能用低精度算的用低精度算不能用低精度更新的坚决保留高精度。所谓“权重主副本”Master Weights就是在内存里始终保留一份 FP32 权重的“真值”。前向传播和反向传播时把这份 FP32 权重临时转换成 FP16 参与运算但优化器更新时则用 FP16 计算出的梯度去更新 FP32 权重而不是直接更新 FP16 权重。为什么非要多此一举因为权重更新是一个长期累积过程。假设一个权重当前是 2048梯度是 0.1学习率是 0.001那么这步更新量是 0.0001。如果权重和梯度都在 FP16 里2048 这个量级下 FP16 根本分辨不了 0.0001 的变化更新等于白做。但如果更新发生在 FP32 权重上FP32 在 2048 附近的精度远高于 FP16这 0.0001 的更新能被准确记录下来。你可以把 FP32 权重主副本理解成一本总账FP16 是草稿纸用来快速做计算FP32 是正式账本每一次最终修订都记在这本账上。如果直接在草稿纸上改账几轮下来账目就乱了。这也是为什么很多框架的混合精度模块会额外占一块显存它不仅存了 FP16 计算副本还保留着 FP32 主副本。这块开销不能省它是训练稳定性的地基。2.2 第二根支柱损失缩放对抗小梯度下溢有了 FP32 主副本下溢问题依然存在反向传播算出来的梯度本身是 FP16如果梯度过小直接在源头就变成 0 了后面拿 FP32 账本也没用。解决办法是损失缩放Loss Scaling。思路非常朴素既然梯度太小会变 0那就先把梯度放大等它过了 FP16 的安全范围再在更新前缩回来。具体怎么做前向传播计算出 loss 后不直接用它做反向传播而是先乘上一个较大的缩放因子比如 1024 或 65536。由于反向传播的梯度是通过链式法则从 loss 逐层传回来的loss 被放大多少倍每一层的梯度也近似被放大多少倍。这样那些原本在 FP16 下会变 0 的小梯度被放大后变得可表示了。等反向传播结束、梯度准备用于更新之前再把这些梯度统一除以缩放因子恢复到原来的数量级然后用它去更新 FP32 主副本。这里有一个关键顺序问题损失缩放只应该在反向传播阶段起作用更新阶段必须还原。如果忘了还原相当于学习率被无意识放大了 scale 倍训练曲线会被破坏。2.3 第三根支柱动态调整与整套数据流固定一个缩放因子行不行不太行。训练不同阶段梯度的数量级差异很大训练初期梯度往往较大固定的大 scale 容易放大成 inf训练后期梯度变得很小固定的小 scale 又压不住下溢。所以现代混合精度实现普遍采用动态损失缩放。它的运行机制很像自动调节水位的阀门训练开始时给一个相对大的 scale如果在反向传播后检测到梯度中出现 inf 或 nan说明 scale 太大就跳过这一轮的参数更新并把 scale 减半如果连续若干步都没有异常就尝试把 scale 翻倍让水位慢慢回升。整套混合精度训练的完整数据流可以概括成六个步骤初始化一份 FP32 权重主副本。前向传播前把 FP32 权重转换成 FP16用 FP16 做矩阵乘等计算对数值敏感的算子仍保持 FP32。计算 loss 后乘上当前缩放因子得到放大后的 loss。用放大后的 loss 做反向传播得到 FP16 的梯度。检查梯度是否出现 inf 或 nan如果有跳过更新并减小 scale如果没有把梯度除以缩放因子还原。用还原后的梯度更新 FP32 权重主副本然后根据最近一段时间的异常频率决定是否调大 scale。这套设计就是混合精度训练的灵魂。很多初学者以为混合精度“就是把模型改成半精度”实际上它是个精密配合的多层系统。理解了这三根支柱后面看任何框架的 AMP 代码都只是细节问题。3. 大模型场景下的典型配置与工程权衡3.1 先算一笔显存账聊完原理我们来算一笔具体的账看看混合精度在大模型场景下到底省了什么、没省什么。假设一个 100B千亿参数的模型用 Adam 优化器训练模型权重FP32 需要 400GBFP16 需要 200GB。Adam 优化器状态每个参数通常维护一阶动量 m 和二阶动量 v 各一份。如果都用 FP32需要 800GB。梯度FP32 需要 400GBFP16 需要 200GB。如果全用 FP32权重、优化器状态、梯度三项合计约 1600GB这还没有算激活值。如果简单粗暴地全部转 FP16三项合计约 800GB但训练稳定性会很差。混合精度的典型配置是权重保留 FP32 主副本400GB优化器状态保留 FP32800GB梯度在反向传播时用 FP16200GB合计约 1400GB。看起来比全 FP32 只省了 200GB但这里没有包含前向传播时需要保存的激活值——在训练大模型时激活值往往能占到显存的大头把它们从 FP32 降到 FP16省下的空间非常可观。更关键的是权重主副本和优化器状态这两大块混合精度并没有帮你砍掉。这就是为什么大模型训练还需要优化器状态分片、激活值重计算等技术。混合精度真正的价值更多体现在“计算吞吐翻倍、通信量减半、临时缓冲显存减半”上而不是把所有存储都变成半精度。3.2 框架视角下的混合精度流水线理解了原理再看实际代码会清晰很多。现在主流框架基本都内置了自动混合精度模块核心 API 通常包含两个角色自动类型转换器和梯度缩放器。自动类型转换器负责在前向传播时自动判断哪些算子可以用低精度哪些必须保留高精度。它不会简单粗暴地把整个模型都转成 FP16而是让矩阵乘法、卷积这类对精度不敏感的算子走 FP16让 LayerNorm、Softmax 这类对动态范围敏感的算子留在 FP32。梯度缩放器则承担了前面提到的损失缩放和动态调整功能。它有一个内部状态记录当前的 scale 值并在每次优化器更新后根据是否有 nan/inf 来调整。一份典型的训练循环伪代码如下scaler GradScaler(init_scale65536, growth_interval2000) optimizer Adam(model.parameters(), lr0.001) for batch in dataloader: optimizer.zero_grad() with autocast(dtypefloat16): loss model(batch) scaler.scale(loss).backward() scaler.unscale_(optimizer) clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update()有几个点值得逐个解释scaler.scale(loss).backward()这一步做的就是把 loss 放大后再反向传播保证反向出来的梯度是放大过的。scaler.unscale_(optimizer)在更新前把梯度还原为什么要有这一步因为你可能需要在更新前做梯度裁剪而裁剪阈值是基于真实梯度量级设定的必须在 unscale 之后裁剪否则阈值会被 scale 污染。scaler.step(optimizer)内部会先检查梯度有没有 inf/nan如果有就跳过这轮更新scaler.update()则根据最近一段时间的表现决定 scale 增大还是减小。3.3 大模型训练里容易忽略的两个精度决策在大模型训练中有两个精度相关的决策容易出问题值得单独提醒。第一个是梯度累积与精度。为了模拟更大的 batch很多训练流程会累积多个微批次的梯度再统一更新。梯度的累积缓冲器最好保留为 FP32。如果直接在 FP16 上累积每次累加都在引入舍入误差累积几十个 batch 后误差可能超过梯度本身的信号。框架默认行为不一定合理建议检查你用的并行框架对梯度缓冲区的精度处理。第二个是分布式通信时的精度。梯度在反向传播后还处于放大状态如果直接在通信模块做梯度规约不同卡之间的梯度先被送到一起再平均这中间的舍入风险会被放大。常见做法是在通信前把梯度还原成 FP32或者在通信时保持 FP32 缓冲区。除非你明确知道通信模块对 FP16 做了特殊处理否则不要盲目用半精度梯度做 AllReduce。这两个点都属于“事情能跑但结果悄悄变坏”的类型问题是训练过程不会立刻报错而是让模型收敛速度悄悄下降。排查时很费劲所以最好一开始就按稳妥的方式配置。4. 现场排障混合精度训练的常见坑与排查表4.1 现象一loss 突然变成 NaN这件事几乎每个做混合精度训练的人都遇到过。常见原因有几个输入数据里本身就带有 NaN损失缩放因子太大导致中间梯度上溢成 inf某个自定义算子或底层 kernel 在 FP16 下数值不稳定。排查可以先看数据把每个 batch 里的标签和特征扫一遍排除坏样本。然后看梯度在反向传播后打印梯度范数如果出现 inf基本可以确认是缩放因子过大的问题。再就是看计算图对比同一个模型在 FP32 下能否稳定训练如果 FP32 正常、FP16 崩溃问题大概率出在某个算子的低精度实现上。一个很实用的临时缓解方案调低初始 scale 值同时调大梯度裁剪的阈值。如果问题消失了说明是损失缩放过大如果问题还在就要往算子和数据方向查。4.2 现象二loss 不降反稳模型几乎不学习这种情况比 NaN 更折磨人因为报错都没有看起来一切正常但曲线就是不动。最常见的原因是梯度下溢大量梯度的绝对值小于 FP16 的最小正常数在反向传播阶段就被截成 0。尤其到了训练中后期梯度整体变小这个问题会越来越突出。你可以做一次简单的检查在 unscale 之后打印梯度张量的分布看看有多少比例的梯度精确等于 0。如果过半梯度都是 0那基本可以断定下溢是元凶。另一个原因是权重更新量被舍入吃掉。虽然我们保留了 FP32 主副本但如果你在代码里手动把主副本也转了 half或者优化器状态没有保持 FP32更新阶段就会重新陷入“大数吃小数”的局面。检查一下模型参数的实际 dtype确认主副本始终是 FP32。实际处理时我会先改成 BF16 试试。BF16 的指数范围和 FP32 一致天然没有下溢问题虽然尾数精度更低但在很多大模型训练场景下表现反而比 FP16 稳定。如果你的加速卡支持 BF16这是一个性价比很高的尝试。4.3 现象三显存没省多少速度反而变慢有些读者会反馈开了混合精度显存好像没降多少速度还变慢了。这不奇怪。显存没降多少可能是因为权重主副本和优化器状态本来就是 FP32混合精度省的主要是激活值。如果你的序列长度和 batch 不算大激活值省下的空间不明显反而多了一份额外的 FP32 主副本净收益自然微薄。要真正压缩显存还得配合优化器状态分片、梯度分片这类手段单纯指望混合精度不现实。速度变慢通常出在“低精度算子没生效”上。混合精度模块只会把支持的算子转成低精度如果某个自定义算子没有对应的 FP16 kernel它会静默退回 FP32再加上自动类型转换器带来的额外判断开销整体速度反而下降。排查方式是做一次逐算子的性能剖析观察关键算子的执行 dtype 和耗时占比。还有一个常见问题是自动转换的范围过大导致频繁在 FP32 和 FP16 之间拷贝数据。对已经调好的工程我会缩小自动转换范围只对计算量最大的矩阵乘法路径开启半精度保留其余部分为 FP32。4.4 踩坑速查表现象常见原因快速验证方式处理建议loss 变 NaN数据含 NaN / scale 过大 / 算子不稳定检查数据样本打印梯度范数清洗数据降低初始 scale尝试强制 FP32 算子loss 纹丝不动梯度下溢为 0 / 主副本精度错误打印梯度分布统计 0 值占比开启损失缩放提高 scale换 BF16显存没明显下降省的是激活值参数和优化器仍占大量统计激活值内存占比结合分片、激活重计算等手段速度不升反降算子静默退回 FP32 / 频繁 cast剖析算子 dtype 和执行时间缩小混合精度范围补齐低精度 kernel梯度裁剪纸失效clip 放在 unscale 之前检查代码顺序先 unscale 再做 clip多卡训练结果不一致通信时梯度精度处理不当对比单卡与多卡 loss 曲线通信前恢复 FP32 梯度或使用 FP32 规约缓冲4.5 一个验证混合精度是否生效的小技巧最后分享一个我常用的验证技巧。训练跑起来之后不用等几十个 step先做两件事第一打印模型权重的 dtype确认主副本是 FP32第二观察梯度缩放器内部的 scale 值。在训练初期scale 值应该会明显变化如果输入数据正常它会在一个范围内上下浮动而不是一条直线。如果 scale 从头到尾完全不变说明你的梯度异常检测逻辑可能根本没生效要么损失缩放没有真正参与反向传播要么梯度一直都正常到不需要调整。另外一个我经常做的对照实验同一份数据和同一个随机种子分别用纯 FP32 和混合精度各训练几步对比每一步的梯度范数。如果两条梯度范数曲线的形状大体一致说明混合精度配置是健康的如果差异超过一个量级就该回查具体算子了。这个对照实验的成本很低却能帮你提前很久发现数值隐患。混合精度训练入了门之后你会发现它不是一个简单的开关而是一整套需要和模型结构、数据规模、并行方式协同设计的系统工程。但正因如此搞清楚它才特别值——它能让你在相同的硬件条件下把训练规模往前推进一大步。
阅读完成 · 觉得有帮助?