模型训练过程中，混合精度训练稳定性解决方案

2023-03-22
山东
本文字数：5473 字
阅读完需：约 18 分钟

以下文章来源于微信公众号：极市平台

作者：Ben Snyder

链接：https://mp.weixin.qq.com/s/ILrr_j-6xJDwTNKnLlTAnQ

本文仅用于学术分享，如有侵权，请联系后台作删文处理

导读

本文主要讨论关于混合精确训练的数值稳定性问题，作者列举了一些解决方案以及数值不稳定性的解决方案，值得收藏学习。

混合精度已经成为训练大型深度学习模型的必要条件，但也带来了许多挑战。将模型参数和梯度转换为较低精度数据类型（如 FP16）可以加快训练速度，但也会带来数值稳定性的问题。使用进行 FP16 训练梯度更容易溢出或不足，导致优化器计算不精确，以及产生累加器超出数据类型范围的等问题。

在这篇文章中，我们将讨论混合精确训练的数值稳定性问题。为了处理数值上的不稳定性，大型训练工作经常会被搁置数天，会导致项目的延期。所以我们可以引入 Tensor Collection Hook 来监控训练期间的梯度条件，这样可以更好地理解模型的内部状态，更快地识别数值不稳定性。在早期训练阶段了解模型的内部状态可以判断模型在后期训练中是否容易出现不稳定是非常好的办法，如果能够在训练的头几个小时就能识别出梯度不稳定性，可以帮助我们提升很大的效率。所以本文提供了一系列值得关注的警告，以及数值不稳定性的补救措施。

混合精度训练

随着深度学习继续向更大的基础模型发展。像 GPT 和 T5 这样的大型语言模型现在主导着 NLP，在 CV 中对比模型(如 CLIP)的泛化效果优于传统的监督模型。特别是 CLIP 的学习文本嵌入意味着它可以执行超过过去 CV 模型能力的零样本和少样本推理，训练这些模型都是一个挑战。

这些大型的模型通常涉及深度 transformers 网络，包括视觉和文本，并且包含数十亿个参数。GPT3 有 1750 亿个参数，CLIP 则是在数百 tb 的图像上进行训练的。模型和数据的大小意味着模型需要在大型 GPU 集群上进行数周甚至数月的训练。为了加速训练减少所需 gpu 的数量，模型通常以混合精度进行训练。

混合精确训练将一些训练操作放在 FP16 中，而不是 FP32。在 FP16 中进行的操作需要更少的内存，并且在现代 gpu 上可以比 FP32 的处理速度快 8 倍。尽管在 FP16 中训练的大多数模型精度较低，但由于过度的参数化它们没有显示出任何的性能下降。

随着英伟达在 Volta 架构中引入 Tensor Cores，低精度浮点加速训练更加快速。因为深度学习模型有很多参数，任何一个参数的确切值通常都不重要。通过用 16 位而不是 32 位来表示数字，可以一次性在 Tensor Core 寄存器中拟合更多参数，增加每个操作的并行性。

但 FP16 的训练是存在挑战性的。因为 FP16 不能表示绝对值大于 65,504 或小于 5.96e-8 的数字。深度学习框架例如如 PyTorch 带有内置工具来处理 FP16 的限制(梯度缩放和自动混合精度)。但即使进行了这些安全检查，由于参数或梯度超出可用范围而导致大型训练工作失败的情况也很常见。深度学习的一些组件在 FP32 中发挥得很好，但是例如 BN 通常需要非常细粒度的调整，在 FP16 的限制下会导致数值不稳定，或者不能产生足够的精度使模型正确收敛。这意味着模型并不能盲目地转换为 FP16。

所以深度学习框架使用自动混合精度(AMP)，它通过一个预先定义的 FP16 训练安全操作列表。AMP 只转换模型中被认为安全的部分，同时将需要更高精度的操作保留在 FP32 中。另外在混合精度训练中模型中通过给一些接近于零梯度（低于 FP16 的最小范围）的损失乘以一定数值来获得更大的梯度，然后在应用优化器更新模型权重时将按比例向下调整来解决梯度过小的问题，这种方法被称为梯度缩放。

下面是 PyTorch 中一个典型的 AMP 训练循环示例。

梯度缩放器 scaler 会将损失乘以一个可变的量。如果在梯度中观察到 nan，则将倍数降低一半，直到 nan 消失，然后在没有出现 nan 的情况下，默认每 2000 步逐渐增加倍数。这样会保持梯度在 FP16 范围内，同时也防止梯度变为零。

训练不稳定的案例

尽管框架都尽了最大的努力，但 PyTorch 和 TensorFlow 中内置的工具都不能阻止在 FP16 中出现的数值不稳定情况。

在 HuggingFace 的 T5 实现中，即使在训练之后模型变体也会产生 INF 值。在非常深的 T5 模型中，注意力值会在层上累积，最终达到 FP16 范围之外，这会导致值无穷大，比如在 BN 层中出现 nan。他们是通过将 INF 值改为在 FP16 的最大值解决了这个问题，并且发现这对推断的影响可以忽略不计。

另一个常见问题是 ADAM 优化器的限制。作为一个小更新，ADAM 使用梯度的第一和第二矩的移动平均来适应模型中每个参数的学习率。

这里 Beta1 和 Beta2 是每个时刻的移动平均参数，通常分别设置为 .9 和 .999。用 beta 参数除以步数的幂消除了更新中的初始偏差。在更新步骤中，向二阶矩参数添加一个小的 epsilon 以避免被零除产生错误。epsilon 的典型默认值是 1e-8。但 FP16 的最小值为 5.96e-8。这意味着如果二阶矩太小，更新将除以零。所以在 PyTorch 中为了训练不会发散，更新将跳过该步骤的更改。但问题仍然存在尤其是在 Beta2=.999 的情况下，任何小于 5.96e-8 的梯度都可能会在较长时间内停止参数的权重更新，优化器会进入不稳定状态。

ADAM 的优点是通过使用这两个矩，可以调整每个参数的学习率。对于较慢的学习参数，可以加快学习速度，而对于快速学习参数，可以减慢学习速度。但如果对多个步骤的梯度计算为零，即使是很小的正值也会导致模型在学习率有时间向下调整之前发散。

另外 PyTorch 目前还一个问题，在使用混合精度时自动将 epsilon 更改为 1e-7，这可以帮助防止梯度移回正值时发散。但是这样做会带来一个新的问题，当我们知道梯度在相同的范围内时，增加ε会降低了优化器适应学习率的能力。所以盲目的增加 epsilon 也不能解决由于零梯度而导致训练停滞的情况。

CLIP 训练中的梯度缩放

为了进一步证明训练中可能出现的不稳定性，我们在 CLIP 图像模型上构建了一系列实验。CLIP 是一种基于对比学习的模型，它通过视觉转换器和描述这些图像的文本嵌入同时学习图像。对比组件试图在每批数据中将图像匹配回原始描述。由于损失是在批次中计算的，在较大批次上的训练已被证明能提供更好的结果。

CLIP 同时训练两个 transformers 模型，一个类似 GPT 的语言模型和一个 ViT 图像模型。两种模型的深度都为梯度增长创造了超越 FP16 限制的机会。OpenClip（arxiv 2212.07143）实现描述了使用 FP16 时的训练不稳定性。

Tensor Collection Hook

为了更好地理解训练期间的内部模型状态，我们开发了一个 Tensor Collection Hook (TCH)。TCH 可以包装一个模型，并定期收集关于权重、梯度、损失、输入、输出和优化器状态的摘要信息。

例如，在这个实验中，我们要找到和记录训练过程中的梯度条件。比如可能想每隔 10 步从每一层收集梯度范数、最小值、最大值、绝对值、平均值和标准差，并在 TensorBoard 中可视化结果。

然后可以用 out_dir 作为--logdir 输入启动 TensorBoard。

实验

为了重现 CLIP 中的训练不稳定性，用于 OpenCLIP 训练 Laion 50 亿图像数据集的一个子集。我们用 TCH 包装模型，定期保存模型梯度、权重和优化器时刻的状态，这样就可以观察到不稳定发生时模型内部发生了什么。

从 vvi - h -14 变体开始，OpenCLIP 作者描述了在训练期间存在稳定性问题。从预训练的检查点开始，将学习率提高到 1-e4，与 CLIP 训练后半段的学习率相似。在训练进行到 300 步时，有意连续引入 10 个难度较大的训练批次。

损失会随着学习率的增加而增加，这是可预期的。当在第 300 步引入难度较大的情况时，损失会有一个小的，但不是很大的增加。该模型发现难度较大的情况，但没有更新这些步骤中的大部分权重，因为 nan 出现在梯度中(在第二个图中显示为三角形)。通过这组难度较大的情况后，梯度降为零。

PyTorch 梯度缩放

这里发生了什么?为什么梯度是零?问题就出在 PyTorch 的梯度缩放。梯度缩放是混合精度训练中的一个重要工具。因为在具有数百万或数十亿个参数的模型中，任何一个参数的梯度都很小，并且通常低于 FP16 的最小范围。

当混合精确训练刚刚提出时，深度学习的科学家发现他们的模型在训练早期通常会按照预期进行训练，但最终会出现分歧。随着训练的进行梯度趋于变小，一些下溢的 FP16 变为零，使训练变得不稳定。

为了解决梯度下溢，早期的技术只是简单地将损失乘以一个固定的量，计算更大的梯度，然后将权重更新调整为相同的固定量(在混合精确训练期间，权重仍然存储在 FP32 中)。但有时这个固定的量仍然不够。而较新的技术，如 PyTorch 的梯度缩放，从一个较大的乘数开始，通常是 65536。但是由于这可能很高，导致大的梯度会溢出 FP16 值，所以梯度缩放器监视将溢出的 nan 梯度。如果观察到 nan，则在这一步跳过权重更新将乘数减半，然后继续下一步。这一直持续到在梯度中没有观察到 nan。如果在 2000 步中梯度缩放器没有检测到 nan，它将尝试使乘数加倍。

在上面的例子中，梯度缩放器完全按照预期工作。我们向它传递一组比预期损失更大的情况，这会产生更大的梯度导致溢出。但问题是现在的乘数很低，较小的梯度正在下降到零，梯度缩放器不监视零梯度只监视 nan。

上面的例子最初看起来可能有些故意的成分，因为我们有意将困难的例子分组。但是经过数天的训练，在大批量的情况下，产生 nan 的异常情况的概率肯定会增加。所以遇到足够多的 nan 将梯度推至零的几率是非常大。其实即使不引入困难的样本，也经常会发现在几千个训练步骤后，梯度始终为零。

产生梯度下溢的模型

为了进一步探索问题何时发生，何时不发生，将 CLIP 与通常在混合精度下训练的较小 CV 模型 YOLOV5 进行了比较。在这两种情况下的训练过程中跟踪了每一层中零梯度的频率。

在前 9000 步的训练中，CLIP 中 5-20%的层显示梯度下溢，而 Yolo 中的层仅显示偶尔下溢。CLIP 中的下溢率也随着时间的推移而增加，使得训练不太稳定。

使用梯度缩放并不能解决这个问题，因为 CLIP 范围内的梯度幅度远远大于 YOLO 范围内的梯度幅度。在 CLIP 的情况下，当梯度缩放器将较大的梯度移到 FP16 的最大值附近时，最小的梯度仍然低于最小值。

如何解决解 CLIP 中的梯度不稳定性

在某些情况下，调整梯度缩放器的参数可以帮助防止下溢。在 CLIP 的情况下，可以尝试修改以一个更大的乘数开始，并缩短增加间隔。

但是我们发现乘数会立即下降以防止溢出，并迫使小梯度回到零。

改进缩放比例的一种解决方案是使其更适应参数范围。比如论文 Adaptive Loss Scaling for Mixed Precision Training 建议按层而不是整个模型执行损失缩放，这样可以防止下溢。而我们的实验表明需要一种更具适应性的方法。由于 CLIP 层内的梯度仍然覆盖整个 FP16 范围，缩放需要适应每个单独的参数以确保训练稳定性。但是这种详细的缩放需要大量内存会减少了训练的批大小。

较新的硬件提供了更有效的解决方案。比如 BFloat16 (BF16) 是另一种 16 位数据类型，它以精度换取更大的范围。FP16 处理 5.96e-8 到 65,504，而 BF16 可以处理 1.17e-38 到 3.39e38，与 FP32 的范围相同。但是 BF16 的精度低于 FP16，会导致某些模型不收敛。但对于大型的 transformers 模型，BF16 并未显示会降低收敛性。

我们运行相同的测试，插入一批困难的观察结果，在 BF16 中，当引入困难的情况时，梯度会出现尖峰，然后返回到常规训练，因为梯度缩放由于范围增加而从未在梯度中观察到 NaN。

对比 FP16 和 BF16 的 CLIP，我们发现 BF16 中只有偶尔的梯度下溢。

在 PyTorch 1.12 及更高版本中，可以通过对 AMP 的一个小更改来启动 BF16。

如果需要更高的精度，可以试试 Tensorfloat32 (TF32)数据类型。TF32 由英伟达在安培 GPU 中引入，是一个 19 位浮点数，增加了 BF16 的额外范围位，同时保留了 FP16 的精度。与 FP16 和 BF16 不同，它被设计成直接取代 FP32，而不是在混合精度下启用。要在 PyTorch 中启用 TF32，在训练开始时添加两行。

这里需要注意的是：在 PyTorch 1.11 之前，TF32 在支持该数据类型的 gpu 上默认启用。从 PyTorch 1.11 开始，它必须手动启用。TF32 的训练速度比 BF16 和 FP16 慢，理论 FLOPS 只有 FP16 的一半，但仍然比 FP32 的训练速度快得多。

如果你用亚马逊的 AWS：BF16 和 TF32 在 P4d、P4de、G5、Trn1 和 DL1 实例上是可用的。

在问题发生之前解决问题

上面的例子说明了如何识别和修复 FP16 范围内的限制。但这些问题往往在训练后期才会出现。在训练早期，模型会产生更高的损失并对异常值不太敏感，就像在 OpenCLIP 训练中发生的那样，在问题出现之前可能需要几天的时间，这回浪费了昂贵的计算时间。

FP16 和 BF16 都有优点和缺点。FP16 的限制会导致不稳定和失速训练。但 BF16 提供的精度较低，收敛性也可能较差。所以我们肯定希望在训练早期识别易受 FP16 不稳定性影响的模型，这样我们就可以在不稳定性发生之前做出明智的决定。所以再次对比那些表现出和没有表现出后续训练不稳定性的模型，可以发现两个趋势。