混合精度训练 · AI Infra 知识库

01 · 一次训练里并非所有数据都同精度

前向激活、乘法输入、累加器、参数、梯度、优化器状态和通信缓冲是不同对象。混合精度指有选择地使用低精度；不能从“矩阵计算用 BF16”推断优化器状态也只有 2 字节，更不能推断所有算子都在低精度执行。

对象 / 运算	为何单独选择精度
大矩阵输入	低精度可调用适配的高吞吐硬件路径。
累加、归一化、损失	可能需要更高精度或数值稳定实现。
参数与优化器状态	更新幅度很小时，低精度存储可能丢失有效变化。
梯度与集合通信	通信 dtype 是另一项策略，不能由前向 dtype 推断。

依据：PyTorch · torch.amp（main） ↗ · Mixed Precision Training · 原始论文 ↗

02 · FP16、BF16 与缩放的角色

FP16 和 BF16 都占 16 bit，但位数分配不同：FP16 指数 5 bit、尾数 10 bit；BF16 指数 8 bit、尾数 7 bit。BF16 范围更大，但在同数量级附近精度更粗。二者都不能简单描述为“精度完全相同的半容量格式”。

FP16 梯度可能下溢。Loss Scaling 在反向前乘 scale，更新前再除回，扩大反向中梯度的量级；遇到非有限值可跳过更新并调整 scale。缩放不修复已经发生的前向溢出，也不能恢复低精度舍入丢失的信息。BF16 通常不需要同样的梯度缩放配方，仍需观察稳定性。

autocast 按算子策略选择 dtype，而 GradScaler 管理梯度缩放。FP8 路径还需要缩放、量化格式与适配内核等配方，不能把示例中的 dtype 换一下就认为完成 FP8 训练。

依据：PyTorch · torch.amp（main） ↗ · Transformer Engine 2.8 · FP8 / FP4 配方 ↗

03 · 先计算容量，再判断速度

自设示例：一个 [4096,4096] 的矩阵有 16,777,216 个元素。FP32 存储约 64 MiB，BF16 / FP16 约 32 MiB。这只是单个矩阵，不是整个训练状态。若原 FP32 参数仍常驻、又生成低精度副本，不能把参数占用直接算成减半。

另一个容量视角：1 亿参数的 FP32 Adam 两个动量状态合计约 800 MB（十进制），不随前向 autocast 自动变成 400 MB。完整账目还包括参数、梯度、激活和工作区。不同配方的主权重与梯度 dtype 会改变结果。

04 · 一个 FP16 训练步骤

PyTorch CUDA / FP16 接入片段 · 假设 model、optimizer、输入和标签已定义

import torch
scaler = torch.amp.GradScaler("cuda")

optimizer.zero_grad(set_to_none=True)
with torch.autocast(device_type="cuda", dtype=torch.float16):
    prediction = model(inputs)
    loss = loss_fn(prediction, 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()

先 unscale 再做梯度裁剪，否则裁剪阈值作用于被放大的梯度。反向一般放在 autocast 上下文外。BF16 方案可以从 BF16 autocast 开始，并按官方示例选择是否需要 scaler；分布式、累积与多优化器情况要匹配各自流程。

依据：PyTorch · AMP 示例（main） ↗

05 · 正确的验收单位是相同质量

与 FP32 基线对照短程 Loss、梯度范数、非有限值和更新是否被跳过。
保持数据、有效批量、学习率和种子可比，做足够长的收敛实验。
分别记录稳态单步耗时、峰值内存和达到同一质量目标的总时间。
检查 dtype 转换和不支持算子的回退，而不是只观察 Tensor dtype。

低精度改变数值误差，因此“每步更快但需要更多步”也可能发生。推理量化有相近的数据表示问题，但量化校准与混合精度训练配方不同，不应合并为一篇泛泛的说明。

06 · NVIDIA 与昇腾怎样对照

NVIDIA 可从 PyTorch AMP 与 Transformer Engine 的低精度路径阅读。昇腾应核对框架适配、CANN 算子和具体型号支持的计算与累加 dtype。CUDA 代码示例不是昇腾可直接运行的代码；公平比较要列出两侧参数、激活、梯度、状态及通信精度。

依据：Transformer Engine 2.8 · FP8 / FP4 配方 ↗ · MindSpeed Core · 特性清单 ↗

07 · 自查与小实验

列出当前训练的五类数据 dtype：参数、激活、梯度、优化器状态、通信。
在小模型上分别跑 FP32、FP16+scaler、BF16，比较 Loss 曲线、峰值显存和跳步次数。
说明为什么只看模型文件大小不能估计训练内存。

流程位置与关联方法

这个方法可能出现在多个位置；计算线包含模型计算与算子实现，执行线包含运行时、内存和通信。

T-P0 · 确定训练目标与策略 T-A1 · 前向计算 T-A3 · 反向传播得到梯度 T-A4 · 优化器更新参数

资料与核查范围

原理依据以原始论文和官方文档为主；算例与时间线是教学推导。文档核查日期：2026-10-10。引用 main / latest 文档时，实际接口以安装版本为准。

代码为接入片段或教学伪代码；本页没有声称完成 NVIDIA / 昇腾硬件实测。