混合精度曾经成为训练大型深度学习模型的必要条件，但也带来了许多挑战。将模型参数和梯度转换为较低精度数据类型（如 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 不稳定性影响的模型，这样咱们就能够在不稳定性产生之前做出理智的决定。所以再次比照那些体现出和没有体现出后续训练不稳定性的模型，能够发现两个趋势。
在 FP16 中训练的 YOLO 模型和在 BF16 中训练的 CLIP 模型都显示出梯度下溢率个别小于 1%，并且随着工夫的推移是稳固的。
在 FP16 中训练的 CLIP 模型在训练的前 1000 步中下溢率为 5 -10%，并随着工夫的推移出现回升趋势。
所以通过应用 TCH 来跟踪梯度下溢率，可能在训练的前 4 - 6 小时内辨认出更高梯度不稳定性的趋势。当察看到这种趋势时能够切换到 BF16。
总结
混合准确训练是训练现有大型根底模型的重要组成部分，但须要特地留神数值稳定性。理解模型的外部状态对于诊断模型何时遇到混合精度数据类型的限度十分重要。通过用一个 TCH 包装模型，能够跟踪参数或梯度是否靠近数值极限，并在不稳固产生之前执行训练更改，从而可能缩小不胜利的训练运行天数。
https://avoid.overfit.cn/post/eb6b99ec76e3464b8f13d70db588b1f6
作者：Ben Snyder