混合精度Mixed Precision Training

神经网络的参数是用浮点精度表示的, 浮点精度的标准是IEEE 754 - Wikipedia,以下是一个FP16数值在内存中存储格式。

随着神经网络模型规模越来越大,如何减少模型占用的内存并且缩短训练时间成为亟需解决的问题,混合精度训练就是其中之一的解决方案,并且几乎不会影响模型训练的效果。

混合精度原理

想象一下,如果模型参数+loss+gradient都是用fp16保存的,fp16的最小值是 6.1 × 1 0 − 5 6.1\times 10^{-5} 6.1×10−5,小于最小值的gradient都会变成0,相当于浪费了一次梯度传播。或许小的gradient并没有很重要,但是积累多次就会变得不可忽略。当前大模型普遍较低的学习率也会加剧这个问题的影响。

因此为了解决这个问题,就需要用更高精度fp32保存一份参数,在正常前向推理和反向传播时都用fp16,计算好得梯度先转换为fp32,再乘以学习率,然后更新到fp32存储得参数上,最终将fp32参数转换成fp16更新模型参数。

整个流程如下如:

这种用fp16和fp32共同训练模型得技术就叫做混合精度训练(MP, Mixed-Precision training),显然MP并不能节省模型加载需要的内存,因为需要多存储一份fp16的参数和梯度,但是用fp16进行模型前向和后向计算,能够减少中间计算值存储需要的内存,这部分内存会随着sequence length和batch size增大而增大,所以只有在这部分中间值占用内存比重较高时才能带来一定的内存节约。

虽然计算时间的影响不大,但是fp16训练时间的确会大大减少,通常是减少1.5~5.5倍。

更多资料:

fastai - Mixed precision training

Understanding Mixed Precision Training | by Jonathan Davis | Towards Data Science

Loss Scale

是不是混合精度训练就完全没有梯度损失了呢,并不是,在反向传播过程中其实已经有部分梯度因为精度原因丢失了(因为正常模型梯度都不会太大,所以我们主要考虑下溢出)。那么如何解决这部分问题呢,就要用到Loss Scale。

原理是将Loss乘以一个比较大的数scale,因为Loss是用fp32存储的,所以scale的选值范围是比较大的。这样因为反向传播链式法则原理,梯度也会放大很多倍,原本下溢出的值也会保存下来。然后在梯度转换成fp32后除以scale,最后更新就与正常混合精度训练一致了。

流程如下:

一般在开始训练时scale会设定成一个比较大的值,如果计算过程中fp16梯度发生上溢出,会跳过当前步的参数更新,并将scale下调。训练log中会输出如下消息:
⚠️ Gradient overflow. Skipping step, loss scaler 0 reducing loss scale to...

相关推荐
肠畔码农18 小时前
深入分布式事务内核:从 2PC/XA 到 Seata AT 模式的架构演进与权衡
分布式·架构
ai小陈18 小时前
PyTorch多GPU分布式训练实战:从单卡脚本迁移到DDP
服务器·人工智能·pytorch·分布式·深度学习·ai·gpu算力
伟大的大威20 小时前
三台 DGX Spark 部署 DeepSeek V4 Flash NVFP4:从零到可用教程
大数据·分布式·spark
阿里云云原生20 小时前
企业级实时数据平台建设:利用存算分离 Kafka 实现低成本、高可靠入湖
分布式·阿里云·云原生·kafka
国科安芯20 小时前
星载CANFD总线通信网络中抗辐射微控制器MCU的失效机理与容错设计研究
网络·人工智能·分布式·单片机·嵌入式硬件·架构·抗辐射加固
zcmodeltech21 小时前
污水处理设备沙盘模型控制系统设计与实现:多单元协同联动方案
网络·分布式·stm32·单片机·嵌入式硬件·交互
roman_日积跬步-终至千里21 小时前
Spark 资源配置与 Shuffle 故障排查生产手册
大数据·分布式·spark
roman_日积跬步-终至千里21 小时前
常驻 Spark 引擎的稳定性风险:从一次 Linkis 任务失败说起
大数据·分布式·spark
JAVA面经实录9171 天前
Kafka面试题标准答案(面试背诵版)
分布式·面试·kafka
hh9501 天前
Agent Plan x DeepSeek Harness:分布式Agent集群调度与负载均衡
运维·分布式·负载均衡·adg·agent plan·adg成都社区