混合精度曾经成为训练大型深度学习模型的必要条件，但也带来了许多挑战。将模型参数和梯度转换为较低精度数据类型（如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