Mixed-Precision Training of Deep Neural Networks | NVIDIA Technical Blog
目录
深度神经网络 (DNN) 在许多领域取得了突破,包括图像处理和理解、语言建模、语言翻译、语音处理、游戏等。为了实现这些结果,DNN 的复杂性一直在增加,这反过来又增加了训练这些网络所需的计算资源。混合精度训练通过使用较低精度的计算(FP16)来降低所需的资源,这具有以下优点。
图 1.大型LSTM?英语语言模型的训练曲线显示了本文中描述的混合精度训练技术的好处。Y 轴是训练损失。不带损耗缩放的混合精度(灰色)在一段时间后会发散,而带损耗缩放的混合精度(绿色)与单精度模型(黑色)匹配。
由于 DNN 训练传统上依赖于 IEEE 单精度格式,因此本文的重点是半精度训练,同时保持单精度实现的网络精度(如图 1 所示)。这种技术称为混合精度训练,因为它同时使用单精度和半精度表示。
半精度浮点格式由 1 个符号位、5 个指数位和 10 个小数位组成。支持的指数值属于 [-24, 15] 范围,这意味着该格式支持 [2-24,65,504]范围。由于这比 [2-149, ~3.4×1038] 范围支持单精度格式,训练某些网络需要额外考虑。本节介绍了成功训练半精度 DNN 的三种技术:将 FP16 产品累积到 FP32 中;损失缩放;以及砝码的 FP32 主副本。借助这些技术, NVIDIA 和百度研究院能够匹配所有经过训练的网络的单精度结果准确性(混合精度训练)。请注意,并非所有网络都需要使用所有这些技术进行训练。
有关如何在各种框架中应用这些技术的详细说明,包括可用的代码示例,请参阅混合精度训练用户指南。
NVIDIA Volta GPU 架构引入了 Tensor Core 指令,该指令将半精度矩阵相乘,将结果累积为单精度或半精度输出。我们发现,累积到单个精度对于获得良好的训练结果至关重要。累积值在写入内存之前转换为半精度。cuDNN 和 CUBLAS 库提供了多种依赖于 Tensor Core 进行算术运算的函数。
训练 DNN 时会遇到四种类型的张量:激活、激活梯度、权重和权重梯度。根据我们的经验,激活、权重和权重梯度落在半精度表示的值大小范围内。然而,对于某些网络,小幅度激活梯度低于半精度范围。例如,考虑图 2 中训练 Multibox SSD 检测网络时遇到的激活梯度直方图,该直方图显示了 log2 刻度上值的百分比。小于 2-24 的值在半精度格式中变为零。
请注意,激活梯度不使用大多数半精度范围,激活梯度往往是幅度小于 1 的小值。因此,我们可以通过将激活梯度乘以比例因子?S?来将它们“移位”到 FP16 表示的范围内。在SSD网络的情况下,将梯度乘以8就足够了。这表明激活梯度值小于 2-27与该网络的训练无关,而保留 [2-27, 2-24) 范围。
图2.以单精度训练 Multibox SSD 检测器网络时记录的激活梯度直方图。Y 轴是对数刻度上所有值的百分比。X 轴是绝对值的对数刻度,也是零的特殊条目。例如,在此培训课程中,66.8% 的值为零,而 4% 的值介于 2 之间-32和 2-30.
确保梯度落入半精度表示的范围内的一种非常有效的方法是将训练损失乘以比例因子。这仅增加了一次乘法,并且通过链式规则,它确保所有梯度都按比例放大(或向上移动),而无需额外费用。损失缩放可确保恢复丢失到零的相关梯度值。在权重更新之前,权重梯度需要按相同的因子?S?缩小。缩减操作可以与权重更新本身融合(导致没有额外的内存访问)或单独执行。有关详细信息,请参阅《混合精度训练用户指南》和《混合精度训练》白皮书。
DNN 训练的每次迭代都会通过添加相应的权重梯度来更新网络权重。权重梯度幅度通常明显小于相应的权重,尤其是在与学习率相乘(或Adam或Adagrad等优化器的自适应计算因子)相乘之后。如果其中一个加法太小而无法产生半精度表示差异,则此幅度差异可能导致不会发生更新(例如,由于指数差大,较小的加法在移位以对齐二进制点后变为零)。
对于以这种方式丢失更新的网络,一个简单的补救措施是以单精度维护和更新权重的主副本。在每次迭代中,都会制作一个主权重的半精度副本,并将其用于正向和反向传播,从而获得性能优势。在权重更新期间,计算出的权重梯度将转换为单精度,并用于更新主副本,并在下一次迭代中重复该过程。因此,我们只在需要的地方将半精度存储与单精度存储混合使用。
上面介绍的三种技术可以组合到每个训练迭代的以下步骤序列中。对传统迭代过程的补充以粗体显示。
AMP(Automatic mixed precision):自动混合精度,该方法在训练网络时将单精度(FP32)与半精度(FP16)结合在一起,它使用FP16即半精度浮点数存储和计算,从而实现节省显存和加快训练速度的目的。
常用的两种实现amp的方式:
FP16和FP32在计算机的不同存储方法:
半精度浮点数 (FP16):?计算机使用 2 字节 (16 位) 存储,表示范围为?[5.9e-8,65504]
单精度浮点数 (FP32)?:计算机使用 4 字节 (32 位) 存储,表示范围为?[1.4e-45,3.4e38],FP32 能够表示的范围要比 FP16 大的多得多。
默认情况下,大多数深度学习框架都采用FP32进行训练。相比与FP32,FP16具有一下优势:
如果我们简单地把模型权重和输入从 FP32 转化成 FP16,虽然可以加快速度,但是模型的精度会被严重影响,原因如下:
Multibox SSD网络训练过程中激活梯度值的直方图。2%的值在[2?34,2?32)范围内,2%的数值在[2–24,2?23)范围内,以及67%的数值为零。
论文里提到下面三个策略:Micikevicius, Paulius, Sharan Narang, et al. “Mixed Precision Training.” ArXiv:1710.03740 [Cs, Stat], February 15, 2018.?https://arxiv.org/abs/1710.03740
在内存中用FP16做储存和乘法从而加速计算,而用FP32做累加避免舍入误差。混合精度训练的策略有效地缓解了舍入误差的问题。
为了解决下溢出的问题,对计算出来的 loss 值进行缩放 (scale),由于链式法则的存在,对 loss 的缩放会作用在每个梯度上,这些梯度会平移到 FP16 的有效范围内。这样就可以用 FP16 存储梯度而又不会溢出了。此外,在进行更新之前,需要先将缩放后的梯度转化为 FP32,再将梯度反缩放(unscale)回去。?
反向传播前:将loss手动增大缩放因子 (loss_scale)倍
反向传播后:将权重梯度缩小缩放因子 (loss_scale)倍,恢复正常值
将模型权重、激活值、梯度等数据用?FP16?来存储,同时维护一份?FP32?的模型权重副本(master-weight)用于更新。前向使用FP16,在反向传播得到 FP16 的梯度以后,将其转化成 FP32 并 unscale,最后更新 FP32 的模型权重。
尽管与单精度训练相比,保持额外的权重副本会使权重的内存需求增加50%,但对整体内存使用的影响要小得多。对于训练来说,由于更大的batch size和每层的输出值被保存以在反向传播过程中重复使用,因此内存消耗主要由这些输出值决定。由于输出值也以半精度格式存储,因此训练深度神经网络的总内存消耗大致减半。
APEX中,用户不需要手动将模型或数据类型转换为.half(),只需要从现有的默认 (FP32) 脚本开始,添加与 Amp API 对应的三行,然后就可以使用混合精度进行训练。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
|
NVIDIA的APEX混合精度库为例,里面提供了多种策略,可以根据不同的场景进行使用:
opt_level:O0(纯FP32),O1和O2是混合精度的不同实现,O3(纯FP16),还有O4,O5使用BFLOAT16
cast_model_type:将模型的参数转换为所需的类型。
patch_torch_functions:patch所有 Torch 函数和 Tensor 方法以执行 Tensor Core 友好的操作,例如 FP16 中的 GEMM 和卷积,以及任何受益于 FP32 中的 FP32 精度的操作。
keep_batchnorm_fp32:将 batchnorm 权重保持在 FP32 ,模型的其余部分是 FP16。
master_weights:保持 FP32 权重。
loss_scale:float值 or "dynamic"(自适应调整损失比例)。动态损失放大(dynamic loss scaling),为了充分利用FP16的范围,缓解舍入误差,将loss*loss_scale。如果产生上溢出,则跳出参数更新,缩小放大倍数使其不溢出。在2000步后再尝试使用大的scale来充分利用FP16的范围。
蓝色为默认值
opt_level | O0 | O1 | O2 | O3 |
---|---|---|---|---|
cast_model_type | torch.float32 | None | torch.float16 | torch.float16 |
patch_torch_functions | False | True | False | False |
patch_torch_functions_type | None | torch.float16 | None | None |
keep_batchnorm_fp32 | None | None(自动设为TRUE) | True | False |
master_weights | False | None | True | False |
loss_scale | 1.0 | "dynamic" | "dynamic" | 1.0 |
纯FP32训练,可作为accuracy的baseline
O1步骤:
inf
和nan
),如果检测到 inf 或 nan? ? ? ? ? ? ? ? ? ? ? i. loss_scale /= 2
? ? ? ? ? ? ? ? ? ? ? ii. 跳过此次更新
以?nn.Linear?为例, 这个模块有两个权重参数?weight?和?bias,输入为?input,前向传播就是调用了?torch.nn.functional.linear(input, weight, bias),对于白名单来说,就是把权重参数?weight?和?bias和input转换为 FP16再进行计算。
黑白名单
lists里有三个文件:functional_overrides.py,tensor_overrides.py,torch_overrides.py分别定义了黑白名单(FP16/FP32的适用情况)
O2步骤:
? ? ? ? ? ? ? ? ? i. loss_scale /= 2
? ? ? ? ? ? ? ? ? ii. 跳过此次更新
纯FP16来当速度的baseline
PyTorch 从 1.6 以后开始支持amp,即torch.cuda.amp module,类似于apex的O1模式:
torch.cuda.amp?给用户提供了较为方便的混合精度训练机制,“方便”体现在两个方面:
1.amp 会自动为算子选择合适的数值精度(FP32、FP16)
该名单在 torch\testing\_internal\autocast_test_lists.py里定义
2.amp 提供了loss_scaling 操作 ,为了防止下溢,将loss乘以一个比例因子,并对缩放后的loss反向传播,然后将梯度除以相同的比例因子
代码格式:
|
nvidia apex官方文档:?Apex (A PyTorch Extension) — Apex 0.1.0 documentation
Micikevicius, Paulius, Sharan Narang, et al. “Mixed Precision Training.” ArXiv:1710.03740 [Cs, Stat], February 15, 2018.?https://arxiv.org/abs/1710.03740
torch.amp文档? https://pytorch.org/docs/stable/amp.html#