在深度学习领域,神经网络的训练性能瓶颈常常出现在 GPU显存的使用上。主要表现为两方面:
- 单卡上可容纳的模型和数据量有限;
- 显存与计算单元之间的带宽和延迟限制了运算速度;
为了解决显卡瓶颈的问题,涌现了不同的解决方法。
模型参数量估计
为了更好地估算模型所要占用的显存,首先需要分析模型训练过程中有哪些部分需要消耗存储空间。在 “ZeRO: Memory Optimizations Toward Training Trillion Parameter Model” 中提出,模型在训练时,主要有两大部分的空间占用。
- 对于大模型来说,主要的空间占用是模型状态,包括优化器状态(eg:Adam优化器的动量和方差)、模型参数和模型参数的梯度;
- 剩余的空间主要被模型训练中间激活值、临时缓冲区和不可用的内存碎片占用,统称为剩余状态。
以 3.5B 大模型为例。①35亿个参数,如果使用 FP16进行存储的话,即 70亿个字节,约7GB左右;②前向传播的激活值和反向传播的梯度大小跟模型参数保持一致,约 7GB;③以 Adam优化器为例,包括三部分,分别为 FP32格式模型参数的备份,FP32的动量和方差,加起来约28GB;因此,从理论上,要微调此模型的话,至少需要 49GB的空间。
1、BF16 半精度浮点数
双精度浮点 Float64;单浮点精度 Float32;半浮点精度 Float16,被广泛应用于model推理
为了进一步降低计算量和显存占用,可以考虑整数int4和int8量化推理
之前深度学习模型的训练通常都采用 Float32(FP32)的精度,而作者发现,使用较低的精度来进行模型的训练也是可行的,并且能够显著提升速度。通过采用混合精度训练,一般可以实现 2~3 倍的速度提升,极大的优化了模型的训练流程。
- FP32 整体长度为4个字节,即32位,其中有8位的指数位宽,23位的尾数精度和1位的符号位,能够表示的数值范围是 1 × 2 − 126 ∼ ( 2 − ϵ ) × 2 127 1\times 2^{-126} \sim (2-\epsilon)\times 2^{127} 1×2−126∼(2−ϵ)×2127 ;
- 在一些不太需要高精度计算的应用中,eg:图像处理和神经网络中,32位的空间其实有一些浪费,因此又出现了新的数据类型,半精度浮点数 FP16,使用16位(2个字节)来存储浮点值,有5位的指数位宽,10位的尾数精度和1位的符号位,能够表示的数值范围是 1 × 2 − 14 ∼ ( 2 − ϵ ) × 2 15 1\times 2^{-14} \sim (2-\epsilon)\times 2^{15} 1×2−14∼(2−ϵ)×215;
格式 | 位数/位 | 指数位宽/位 | 尾数精度/位 | 符号位/位 | 数值范围 |
---|---|---|---|---|---|
FP32 | 32 | 8 | 23 | 1 | 1 × 2 − 126 ∼ ( 2 − ϵ ) × 2 127 1\times 2^{-126} \sim (2-\epsilon)\times 2^{127} 1×2−126∼(2−ϵ)×2127 |
FP16 | 16 | 5 | 10 | 1 | 1 × 2 − 14 ∼ ( 2 − ϵ ) × 2 15 1\times 2^{-14} \sim (2-\epsilon)\times 2^{15} 1×2−14∼(2−ϵ)×215 |
BP32 | 16 | 8 | 7 | 1 | 1 × 2 − 126 ∼ ( 2 − ϵ ) × 2 127 1\times 2^{-126} \sim (2-\epsilon)\times 2^{127} 1×2−126∼(2−ϵ)×2127 |
混合精度训练,即在模型训练时同时采用 FP32 和 FP16 两种精度。在实践过程中,研究人员发现在大语言模型的训练中直接使用 FP16会有一些问题,在训练过程中 loss 会非常不稳定,因此使用 FP16 训练大模型非常困难。问题在于 FP16的指数位宽只有 5位,能表示的最大整数为 65504,一旦权重超过这个值就会发生溢出,因此只能进行较小数的乘法,eg:可以计算 250 × 250 = 62500 250\times250=62500 250×250=62500,但如果计算 255 × 255 = 65025 255\times 255=65025 255×255=65025 就会溢出,这是导致训练出现问题的主要原因。这也意味着模型权重必须保持很小。一种成为损失缩放的技术可以缓解这个问题,但是当模型变得非常大时,FP16 较小的数值范围依旧是一个问题。
- 为了更好地解决 FP16的问题,谷歌开发了一种新的浮点数格式 BF16(Brain Floating Point, 2个字节),用于降低存储需求,提高机器学习算法的计算速度。BF16 的指数位宽为8位,于 FP32相同,尾数精度采用7位,因此当使用 BF16时,精度非常差。然而,在训练模型时一般采用随机梯度下降法及其变体,其过程像蹒跚而行,即使某一步没有找到最优方向也没关系,模型会在后续调整纠正。
将模型参数类型从 FP16换为 BF16,训练的大模型 loss值的下降也会变得更加稳定。
这种低精度和 混合精度训练的方法逐渐被广泛接受和应用,深度学习框架、GPU以及 神经网络加速器的设计也因此受到了深渊的影响。可以说,混合精度训练的提出,对深度学习领域起到了关键的推动作用,有效地解决了 GPU显存的使用问题,提升了模型训练的效率。
2、混合精度训练
paper:Mixed Precision Training
- 维护一个权重的单精度副本,在每个优化器步骤后累计梯度(对于前向和反向传播,此副本四舍五入到半精度);
- 提出了损失缩放来保持小幅度的梯度值;
- 使用半精度算法,该算法累积为单精度输出,在存储到内存之前将其转化为半精度;
FP32 为主副本权重,
在混合精度训练时,权重、激活函数、梯度被保存为 FP16,为了与 FP32模型的精度相匹配,在optimizer step时,维持 FP32权重为主线,并使用权重梯度进行更新。在每次迭代时,主权重的 FP16副本用于前向和反向传播,将 FP32训练所需的存储和带宽减半,如上图所示。
虽然对 FP32住权重的需求并不普遍,但许多模型还需要 FP32的两个可能原因是:
- 权重更新变得太小,无法在 FP16中表示,任何大小小于 2 − 24 2^{-24} 2−24 的值在 FP16中都将变为 零,当与学习率相乘时,这些小值梯度在优化器中都会变为零,并对模型的准确性产生不利影响。使用单精度进行更新可以解决这一问题;
- 权重值 与 权重更新的比例非常大。在这种情况下,即使权重更新可以在 FP16中表示,当加法操作将其右移以使二进制点与权重对齐时,它仍然可能变为零。当归一化权重值的幅度比权重更新的新幅度大 至少2048倍,就会发生这种情况。由于 FP16有10位尾数,隐式位必须右移11位或更多位置,才能潜在地创建一个零。在比例大于2048时,隐式比特将右移12位或更多位。这将导致权重更新变得无法恢复的零。更大的比例将导致非标注化数字的效果。同样,可以通过计算 FP32中的更新来抵消这种影响。
图2-a所示,在 FP16前后传播更新 FP32权重主线时,匹配FP32训练结果,而更新FP16权重会导致 80%的相对精度损失。
由于更大的 batch-size 和每层的激活被保存以在反向传播过程中重复使用,因此训练内存消耗主要由激活决定。由于激活也以半精度格式存储,因此训练深度神经网络的整体内存消耗大约减半。
2.1 损失缩放
FP16指数偏差将归一化指数的范围集中到 [-14,15],而实践中的梯度值往往由小幅度(负指数)主导,如图3所示,显示了Multibox SSD模型的 FP32训练期间在所有层上的急活梯度值的直方图,FP16 可表示范围的大部分未必使用,而许多值低于最小可表示范围 变为指数为0。放大梯度将使它们占据更多的可表示范围,并保留否则会丢失为0的值。当梯度未被缩放时,这个特定的网络会发散,但将其压缩8倍(指数为3)就足以匹配 FP32训练所达到的精度。这说明激活
2.2运算精度
神经网络模型分为三类:vector dot-products, reductions, point-wise operations。当涉及到降低精度的算法时,这些类别受益于不同的处理,为了保持模型的准确性,发现一些模型要求 FP16矢量点积将部分累加成 FP32,在写入内存之前将其转换为 FP16。如果 FP32中没有这种积累,一些 FP16模型与基线模型的精度不匹配。
之前的GPU只支持 FP16 乘/加法运算,而 NVIDIA Volta GPU引入了 Tensor Core,可以将 FP16输入矩阵相乘,并将乘积累加到 FP16或 FP32输出中。
FP32中应进行大幅缩减(向量各元素之和)。在累积统计数据和 softmax层时,这种建好主要出现在 batch-normalization 层中。在两种层类型种,仍然从内存中读取和写入FP16张量,在 FP32中执行算术运算,这并没有减缓训练过程,因为这些层的内存带宽有限,对算术速度不敏感。
逐点操作,如非线性和逐像素矩阵运算,是内存带宽有限,由于算术精度不影响这些运算的速度,因此可以使用FP16 或 FP32.