混合精度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...

相关推荐
hzulwy31 分钟前
Kafka基础理论
分布式·kafka
明达智控技术2 小时前
MR30分布式IO在全自动中药煎药机中的应用
分布式·物联网·自动化
jakeswang2 小时前
细说分布式ID
分布式
失散133 小时前
分布式专题——1.2 Redis7核心数据结构
java·数据结构·redis·分布式·架构
王中阳Go4 小时前
头一次见问这么多kafka的问题
分布式·kafka
boonya6 小时前
Kafka核心原理与常见面试问题解析
分布式·面试·kafka
KIDAKN7 小时前
RabbitMQ 重试机制 和 TTL
分布式·rabbitmq
JAVA学习通7 小时前
【RabbitMQ】----初识 RabbitMQ
分布式·rabbitmq
项目題供诗11 小时前
Hadoop(八)
大数据·hadoop·分布式
在未来等你12 小时前
Kafka面试精讲 Day 7:消息序列化与压缩策略
大数据·分布式·面试·kafka·消息队列