混合精度Mixed Precision Training

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

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

混合精度原理

想象一下,如果模型参数+loss+gradient都是用fp16保存的,fp16的最小值是 <math xmlns="http://www.w3.org/1998/Math/MathML"> 6.1 × 1 0 − 5 6.1\times 10^{-5} </math>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 分钟前
Java-72 深入浅出 RPC Dubbo 上手 生产者模块详解
java·spring boot·分布式·后端·rpc·dubbo·nio
橘子在努力4 小时前
【橘子分布式】Thrift RPC(理论篇)
分布式·网络协议·rpc
lifallen6 小时前
Kafka 时间轮深度解析:如何O(1)处理定时任务
java·数据结构·分布式·后端·算法·kafka
沈健_算法小生8 小时前
基于SpringBoot3集成Kafka集群
分布式·kafka·linq
Swift社区9 小时前
ELK、Loki、Kafka 三种日志告警联动方案全解析(附实战 Demo)
分布式·elk·kafka
chanalbert17 小时前
Nacos 技术研究文档(基于 Nacos 3)
spring boot·分布式·spring cloud
线条119 小时前
Spark 单机模式安装与测试全攻略
大数据·分布式·spark
C182981825751 天前
分布式ID 与自增区别
分布式
码字的字节1 天前
深入解析Hadoop架构设计:原理、组件与应用
大数据·hadoop·分布式·hadoop架构设计
悟能不能悟1 天前
Dubbo跨越分布式事务的最终一致性陷阱
分布式·wpf·dubbo