混合精度已经成为训练大型深度学习模型的必要条件，但同时也带来了诸多挑战。将模型参数和梯度转换为精度较低的数据类型（如FP16）可以加快训练速度，但也会带来数值稳定性问题。使用FP16训练梯度更容易溢出或不足，导致优化器计算不准确，以及累加器超出数据类型范围等问题。在这篇文章中，我们讨论了混合精度训练的数值稳定性。为了应对数值不稳定性，大型培训工作通常会被搁置数天，从而导致项目延误。所以我们可以引入TensorCollectionHook来监控训练过程中的梯度情况，这样可以更好地了解模型的内部状态，更快地识别数值不稳定性。在早期训练阶段了解模型的内部状态，可以很好地判断模型在后期训练中是否容易出现不稳定。如果我们能够在训练的前几个小时识别出梯度不稳定性，它可以帮助我们提高很多。效率高。因此，本文提供了一系列需要注意的注意事项，以及数值不稳定性的补救措施。混合精度训练继续进行深度学习，朝着更大的基础模型发展。GPT和T5等大型语言模型现在在NLP中占据主导地位，而在CV中，CLIP等对比模型的泛化能力优于传统的监督模型。特别是，CLIP的学习文本嵌入意味着它可以执行超出过去CV模型能力的零样本和少样本推理，这两种模型的训练都是一项挑战。这些大型模型通常涉及深度转换器网络，包括视觉和文本，并包含数十亿个参数。GPT3有1750亿个参数，CLIP是在数百TB的图像上训练的。模型和数据的大小意味着模型需要在大型GPU集群上进行数周甚至数月的训练。为了加快训练速度并减少所需的GPU数量，模型通常以混合精度进行训练。混合精度训练将一些训练操作放在FP16而不是FP32中。在FP16中执行的操作需要更少的内存，并且在现代GPU上的处理速度比FP32快8倍。虽然大多数在FP16中训练的模型精度较低，但它们并没有表现出由于过度参数化而导致的任何性能下降。随着NVIDIA在Volta架构中引入TensorCores，低精度浮点加速训练甚至更快。因为深度学习模型有很多参数，任何一个参数的准确值通常并不重要。通过用16位而不是32位表示数字，可以一次在TensorCore寄存器中容纳更多参数，从而增加每个操作的并行性。但是FP16的训练很有挑战性。因为FP16不能表示绝对值大于65504或者小于5.96e-8的数。PyTorch等深度学习框架带有内置工具来处理FP16的局限性（梯度缩放和自动混合精度）。但即使进行了这些安全检查，由于参数或梯度超出可用范围，大型训练作业失败的情况并不少见。深度学习的某些组件在FP32中运行良好，但例如BN通常需要非常细粒度的调整，这可能导致FP16极限下的数值不稳定，或者无法产生足够的精度使模型正确收敛。这意味着模型不能盲目地转换为FP16。因此，深度学习框架使用自动混合精度(AMP)，它传递预定义的FP16训练安全操作列表。AMP仅转换模型中被认为安全的部分，同时在FP32中保留需要更高精度的操作。另外，在混合精度训练模型中，通过将一些接近于零梯度（FP16的最小范围以下）的损失乘以一定的值得到一个较大的梯度，然后在应用优化器更新模型权重时，它会按比例向下调整以解决梯度太小的问题，这种方法称为梯度缩放。下面是PyTorch中典型AMP训练循环的示例。梯度缩放器会将损失乘以可变数量。如果在梯度中观察到nans，则multiplier减少一半，直到nans消失，然后在没有nans的情况下，默认每2000步逐渐增加multiplier。这将使梯度保持在FP16范围内，同时防止梯度变为零。训练不稳定的情况尽管框架尽了最大努力，但PyTorch和TensorFlow中内置的工具都无法防止FP16中的数值不稳定。在HuggingFace的T5实现中，模型变体即使在训练后也会产生INF值。在非常深的T5模型中，attention值会越层累加，最终会到达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同时训练两个transformer模型，一个类GPT语言模型和一个ViT图像模型。两种模型的深度都为超过FP16限制的梯度增长创造了机会。OpenClip(arxiv2212.07143)实现描述了使用FP16时的训练不稳定性。TensorCollectionHook为了更好地了解训练期间的内部模型状态，我们开发了TensorCollectionHook(TCH)。TCH可以包装模型并定期收集有关权重、梯度、损失、输入、输出和优化器状态的摘要信息。例如，在本实验中，我们将查找并记录训练过程中的梯度条件。例如，您可能希望每10步从每一层收集梯度范数、最小值、最大值、绝对值、平均值和标准差，并在TensorBoard中可视化结果。然后可以使用out_dir作为--logdir输入启动TensorBoard。实验为了重现CLIP中的训练不稳定性，Laion50亿图像数据集的一个子集用于OpenCLIP训练。我们用TCH包装模型，TCH会定期保存模型梯度、权重和优化器矩的状态，以便在出现不稳定时观察模型内部发生的情况。从VVI-H-14变体开始，OpenCLIP作者描述了训练期间的稳定性问题。从预训练的检查点开始，将学习率提高到1-e4，类似于CLIP训练后半部分的学习率。当训练进行到300步时，有意连续引入10个高难度训练批次。损失随着学习率的增加而增加，这是可以预料的。当在第300步引入更难的情况时，损失会小幅增加，但不会很大。该模型找到了困难的情况，但没有更新这些步骤中的大部分权重，因为nans出现在梯度中（在第二个图中显示为三角形）。通过这组困难的情况后，梯度下降到零。PyTorch梯度缩放在这里发生了什么？为什么梯度为零？问题是PyTorch的梯度缩放。梯度缩放是混合精度训练中的重要工具。因为在具有数百万或数十亿参数的模型中，任何一个参数的梯度都很小，并且通常低于FP16的最小范围。当混合精度训练首次被提出时，深度学习科学家发现他们的模型往往在训练初期按预期训练，但最终却出现了分歧。随着训练的进行，梯度趋于变小，一些下溢的FP16变为零，使训练不稳定。为了解决梯度下溢，早期的技术只是将损失乘以一个固定的量，计算出一个更大的梯度，然后将权重更新调整为相同的固定量（在混合精度训练期间，权重仍然存储在FP32中）。但有时这个固定数额仍然不够。而较新的技术，例如PyTorch的梯度缩放，从较大的乘数开始，通常为65536。但由于这可能很高，导致大梯度溢出FP16值，梯度缩放器会监视将溢出的nan梯度。如果观察到nan，则跳过此步骤的权重更新以将乘数减半并继续下一步。这一直持续到在梯度中没有观察到nans。如果梯度缩放器在2000步内未检测到nans，它将尝试将乘数加倍。在上面的示例中，渐变缩放器完全按照预期工作。我们将一组损失大于预期的情况传递给它，这会产生更大的梯度，导致溢出。但问题是现在乘数很低，较小的梯度下降到零，梯度缩放器不只监控零梯度nan。上面的示例乍一看似乎是故意的，因为我们有意将困难的示例分组。但是经过几天的训练，在大batch的情况下，产生nan异常的概率肯定会增加。所以遇到足够多的nans将梯度推到零的概率非常高。事实上，即使不引入困难样本，也经常会发现经过几千个训练步骤后，梯度始终为零。产生梯度下溢的模型为了进一步探索问题何时发生和何时不发生，将CLIP与YOLOV5进行了比较，YOLOV5是一种较小的CV模型，通常以混合精度进行训练。在这两种情况下，训练期间都会跟踪每一层中零梯度的频率。在训练的前9000步中，CLIP中有5-20%的层出现梯度下溢，而Yolo中的层仅偶尔出现下溢。CLIP中的下溢率也会随着时间的推移而增加，从而使训练不太稳定。使用梯度缩放并不能解决这个问题，因为CLIP范围内的梯度幅度远大于YOLO范围内的梯度幅度。在CLIP的情况下，虽然梯度缩放器将较大的梯度移动到FP16的最大值附近，但最小的梯度仍然低于最小值。如何修复CLIP中的梯度不稳定性在某些情况下，调整梯度缩放器的参数可以帮助防止下溢。在CLIP的情况下，可以尝试修改以从更大的乘数开始并缩短增量间隔。但是我们发现乘数会立即下降以防止溢出并强制小梯度返回零。改进缩放比例的一种解决方案是使其更适应参数范围。例如，论文AdaptiveLossScalingforMixedPrecisionTraining建议按层而不是整个模型执行损失缩放，这样可以防止下溢。相反，我们的实验表明需要一种更具适应性的方法。由于CLIP层内的梯度仍然覆盖整个FP16范围，缩放需要适应每个单独的参数以确保训练稳定性。但是这种详细的缩放需要大量内存并减少了训练批次的大小。更新的硬件提供更高效的解决方案。例如BFloat16(BF16)是另一种16位数据类型，它以精度换取更大的范围。FP16处理5.96e-8到65,504，而BF16可以处理1.17e-38到3.39e38，与FP32的范围相同。但是BF16的准确率低于FP16，会导致部分模型不收敛。但对于大型变压器，BF16似乎不会降低收敛性。我们运行相同的测试，插入一批硬观察，在BF16中引入硬情况时梯度尖峰，然后返回常规训练作为梯度缩放，因为在梯度NaN中从未观察到范围增加。比较FP16和BF16的CLIP，我们发现BF16中仅偶尔出现梯度下溢。在PyTorch1.12及更高版本中，可以通过对AMP进行少量更改来启用BF16。如果您需要更高的精度，请尝试使用Tensorfloat32(TF32)数据类型。TF32由Nvidia在AmpereGPU中引入，是一个19位浮点数，它增加了BF16的额外范围位，同时保留了FP16的精度。与FP16和BF16不同，它旨在直接替代FP32，而不是在混合精度中启用它。要在PyTorch中启用TF32，请在训练开始时添加两行。此处注意：在PyTorch1.11之前，TF32在支持此数据类型的GPU上默认启用。从PyTorch1.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%并且随着时间的推移保持稳定。在训练的前1000步中，在FP16中训练的CLIP模型下溢5-10%，并且随着时间的推移趋于增加。因此，通过使用TCH跟踪梯度下溢率，我们能够在训练的前4-6小时内识别出更高梯度不稳定性的趋势。当观察到这种趋势时切换到BF16。总结混合精度训练是训练现有大规模基础模型的重要组成部分，但需要特别注意数值稳定性。了解模型的内部状态对于诊断模型何时遇到混合精度数据类型的限制非常重要。通过用TCH包装模型，可以跟踪参数或梯度是否接近数值限制，并在不稳定发生之前执行训练更改，从而有可能减少训练运行失败的天数。
