前言

本文属于搬运内容，原作者：Ben Snyder （Salute！），原文链接在文章末尾处。

混合精度训练的重要性与挑战

混合精度训练已成为训练大型深度学习模型的关键技术，但也带来了一些挑战。将模型参数和梯度转换为更低精度的数据类型（如FP16）可以加速训练，但也会引入数值稳定性问题。FP16训练更容易出现梯度下溢或上溢、优化器计算精度不足以及累加器超出数据类型范围的情况。

在这篇文章中，我们将讨论混合精度训练中的数值稳定性挑战。许多大型训练任务常常因为数值不稳定而停滞数天。我们引入了一种张量收集钩子（Tensor Collection Hook）来监控训练过程中梯度的状态，并展示了通过更好地理解模型内部状态，可以更早地识别数值不稳定性。

实际上，在训练早期了解模型的内部状态可以帮助预测模型在后续训练中是否会出现不稳定性。在我们的实验中，我们能够在训练的最初几个小时内识别出梯度不稳定性，而这些不稳定性通常需要几天后才会影响收敛。基于这些实验，我们提供了一系列需要注意的警告信号以及解决数值不稳定性的措施。

混合精度训练

深度学习正逐步迈向更大的基础模型（Foundation Models）。像GPT和T5这样的超大语言模型已在自然语言处理（NLP）领域占据主导地位。而在计算机视觉（CV）领域，对比学习模型（如CLIP）相比传统的监督模型显示出了更强的泛化能力。特别是，CLIP通过学习的文本嵌入，能够实现零样本和少样本推理，超越了以往CV模型的能力。然而，训练基础模型充满了挑战。

基础模型通常涉及用于视觉和文本的深层变换器网络，总参数量达到数十亿，且需要处理海量数据。ChatGPT拥有1750亿个参数，而CLIP则是在数百TB的图像数据上训练的。这些模型和数据的规模意味着基础模型需要数周甚至数月的时间，在大型GPU集群上进行训练。为了加速训练并减少所需GPU的数量，基础模型通常采用混合精度训练。

混合精度训练将部分训练操作置于FP16，而非FP32中。FP16操作需要的内存更少，并且在现代GPU上处理速度最多可达到FP32的8倍。尽管精度较低，但由于模型的高度过参数化，大多数使用FP16训练的模型并未表现出任何可测的性能下降。

加速低精度浮点数训练的浪潮始于Nvidia Volta架构中Tensor Cores的引入。由于深度学习模型参数数量庞大，单个参数的精确值往往并不重要。因此，通过将数字用16位表示而非32位，我们可以在Tensor Core寄存器中一次容纳更多参数，从而增加每次操作的并行性。

混合精度训练的挑战

然而，FP16训练也带来了挑战。FP16无法表示绝对值大于65,504或小于5.96e-8的数字。深度学习框架（如PyTorch）提供了一些内置工具来处理FP16的限制（例如梯度缩放器和自动混合精度）。但即使有这些安全措施，大型训练任务仍可能因为参数或梯度超出可用范围而失败。一些深度学习操作在FP32中表现得更好，例如，批归一化（Batch Normalization）通常需要在FP16范围的极限附近进行非常细致的调整。因此，批归一化可能会受到数值不稳定的影响，或者无法为模型的正确收敛提供足够的精度。这意味着模型不能简单地转换为FP16并期望与FP32训练结果一致。

解决方案是深度学习框架使用自动混合精度（AMP）。AMP是一份预定义的操作列表，标识哪些操作在FP16训练中是安全的。AMP会将模型中被认为安全的部分转换为FP16，同时保留需要更高精度的操作在FP32中。此外，在混合精度训练中，缩放梯度也是一个良好的实践。一些梯度可能接近于零，低于FP16的最小范围。解决方案是将损失乘以一定倍数，计算较大的梯度，然后在优化器更新模型权重时再将缩放倍数调整回来。以下是一个典型的PyTorch AMP训练循环示例：

梯度缩放器的作用

梯度缩放器会将损失乘以一个可变的倍数。如果在梯度中观察到NaN，缩放器会将倍数减半，直到NaN消失；随后，如果NaN未出现，每隔2000步逐渐增加倍数。这种方法试图将梯度保持在FP16的范围内，同时避免梯度变为零。

不稳定性的案例

尽管PyTorch和TensorFlow内置了许多工具来处理FP16的数值不稳定性问题，但它们仍无法完全避免所有不稳定情况的发生。

在HuggingFace的T5实现中，模型较大的变体即使在训练后也会产生无穷值。调查发现，在非常深的T5模型中，注意力值会在层之间累积，最终超出FP16的表示范围，从而在归一化层中产生无穷值和NaN。他们通过将无穷值限制在FP16的最大值来解决了这个问题，并发现这种处理对推理的影响可以忽略不计。

ADAM优化器的限制

另一个常见问题是与ADAM优化器的限制有关。简单回顾一下，ADAM使用梯度的一阶和二阶矩的移动平均值，为模型中的每个参数调整学习率。

来源：InsideAIML

Beta1和Beta2分别是每个矩的移动平均参数，通常设置为0.9和0.999。通过除以Beta参数的幂次，可以消除更新中的初始偏差。在更新步骤中，我们向二阶矩参数添加一个小的epsilon值，以避免被零除。默认的epsilon值通常为1e-8。但请注意，FP16的最小值为5.96e-8。

这意味着如果二阶矩值太小，更新将会出现除以零的情况。在PyTorch中，这种情况会跳过该步的更新以避免训练发散。然而问题仍然存在，尤其是在Beta2=0.999的情况下，任何小于5.96e-8的梯度可能会导致该参数的权重更新在较长时间内停止。此外，优化器会进入不稳定状态。

ADAM的优势在于，通过使用两个矩，我们可以为每个参数自适应地调整学习率。对于学习较慢的参数，我们加快学习率；对于学习较快的参数，我们降低学习率。然而，如果梯度在多次步骤中被计算为零，即使是一个小的正值也可能在学习率来不及调整下降之前导致模型发散。

CLIP训练中的梯度缩放

为了进一步演示训练中可能出现不稳定性的地方，我们对CLIP图像模型进行了一系列实验。CLIP是一种基于对比学习的模型，能够通过视觉Transformer学习图像，同时学习描述这些图像的文本嵌入。对比学习的目标是在每个数据批次中将图像与其原始描述匹配。由于损失是在整个批次上计算的，因此使用较大的批次进行训练已被证明能获得更好的效果。

来源：OpenAI

CLIP同时训练两个Transformer模型：一个类似GPT的语言模型和一个ViT图像模型。这两个模型的深度为梯度提供了超出FP16限制的机会。事实上，OpenCLIP的实现报告了使用FP16时的训练不稳定性问题，模型在训练到一半时出现了NaN。

Tensor Collection Hook（张量收集钩子）

为了更好地了解训练过程中模型的内部状态，我们开发了一个Tensor Collection Hook (TCH)。TCH包装模型，并定期收集权重、梯度、损失、输入、输出以及优化器状态的汇总信息。

例如，在这次实验中，我们关注训练过程中梯度的状态。我们可以每10步收集一次每层的梯度范数、最小值、最大值、绝对值、均值和标准差，并将结果可视化到TensorBoard中。

我们可以通过设置out_dir为--logdir的输入参数启动TensorBoard。

实验设置

为了重现CLIP中的训练不稳定性，我们从OpenCLIP的实现入手，并使用训练OpenCLIP时用到的Laion 50亿图像数据集的一个子集。我们用一个张量收集钩子包装模型，定期保存模型的梯度、权重以及优化器矩的状态，以便观察不稳定性发生时模型内部的情况。所有实验都在Amazon SageMaker P4d.24xlarge实例上进行。

我们从Vit-H-14变体开始，这是OpenCLIP作者描述在训练中遇到稳定性问题的版本。从预训练检查点开始，我们将学习率升温到1e-4，与CLIP训练后半部分的学习率相似。在训练的第300步，我们故意连续引入10个困难的训练批次。

随着学习率的增加，损失也会增加，这是预期的。当在第300步引入困难的案例时，损失出现了小幅但不严重的增长。模型能够处理这些困难案例，但由于梯度中出现了NaN（在第二个图中以三角形表示），大多数步骤中模型未能更新权重。在通过这组困难案例后，梯度下降到零。

目前，PyTorch存在一个未解决的问题，即在使用混合精度时自动将epsilon调整为1e-7。这可以帮助防止梯度回归到正值时的发散问题。但这样做也带来了新的问题：当我们知道梯度处于同一范围时，增加epsilon实际上降低了优化器调整学习率的能力。此外，增加epsilon也无法解决因梯度为零导致训练停滞的情况。

PyTorch 梯度缩放器

这是怎么回事？为什么梯度会变为零？问题出在PyTorch的梯度缩放器上。梯度缩放是混合精度训练中的重要工具。在拥有数百万甚至数十亿参数的模型中，任何单个参数的梯度通常都很小，往往低于FP16的最小表示范围。

当混合精度训练首次流行起来时，深度学习科学家发现他们的模型在训练初期通常表现正常，但最终会发散。随着训练的进行，梯度往往会变得更小，其中一些会下溢到FP16的零值，导致训练变得不稳定。

来源：Nvidia

为了解决梯度下溢，早期的方法是简单地将损失乘以一个固定值，从而计算出更大的梯度，然后通过相同的固定值调整权重更新（在混合精度训练中，权重仍然以FP32存储）。但有时固定值仍不足够。较新的方法（如PyTorch的梯度缩放器）从一个较大的乘数值开始，通常是65536。然而，由于这个值可能过高，较大的梯度可能会导致FP16溢出，梯度缩放器会监测梯度中的NaN（表明溢出）。如果检测到NaN，该步的权重更新会被跳过，乘数会减半，接着进行下一步。这个过程会持续，直到梯度中不再检测到NaN。如果梯度缩放器在2000步内未检测到NaN，它将尝试将乘数加倍。

在上述案例中，梯度缩放器正按照预期工作。我们向模型提供了一组生成超出预期损失的案例，这些案例产生了更大的梯度，导致溢出。然而问题在于，此时乘数变得非常小，以至于较小的梯度下降为零。而梯度缩放器并不会监测零梯度，仅监测NaN。

上述示例可能看起来有些人为，因为我们故意将困难的样本分组。但在几天的训练中，随着批量大小的增加，异常案例产生NaN的概率也会增加。结果是，出现足够多的NaN将梯度压缩到零的可能性也显著增加。事实上，即使不引入困难样本，我们也常常发现梯度在几千次训练步骤后停留在零值。

产生梯度下溢的模型

为了进一步探讨问题何时发生以及何时不会发生，我们将CLIP与另一种常见的混合精度训练模型YOLOv5进行了比较。在这两种情况下，我们跟踪了训练过程中每一层梯度为零的频率。

在训练的前9000步中，CLIP的5%-20%的层出现了梯度下溢，而YOLO的层仅偶尔出现下溢。此外，CLIP的下溢率随着时间的推移而增加，使得训练变得不稳定。

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

解决CLIP中的梯度不稳定问题

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

然而，我们发现，乘数会立即下降以防止溢出，从而将小梯度再次压缩为零。

一种可能的解决方案是使缩放更加适应参数范围。论文 Adaptive Loss Scaling for Mixed Precision Training 建议对每一层而不是整个模型进行损失缩放，以防止下溢。

然而，我们的实验表明，甚至需要更具适应性的方法。由于CLIP中的每一层梯度仍然覆盖整个FP16范围，缩放需要适应每个参数，以确保训练的稳定性。然而，这种细化的缩放需要大量的内存，从而减少训练批量的大小。

使用BFloat16解决梯度不稳定问题

更新的训练硬件提供了一种更有效的解决方案。BFloat16 (BF16) 是一种替代的16位数据类型，用更大的范围取代了一部分精度。FP16的范围是5.96e-8到65,504，而BF16的范围是1.17e-38到3.39e38，与FP32相同。尽管BF16的精度比FP16更低，但对于大规模Transformer模型，BF16并未显示出会降低收敛性的证据。

运行相同的测试，在BF16中插入一批困难的样本时，梯度会在引入困难样本时出现峰值，然后恢复正常训练。由于BF16的更大范围，梯度缩放器没有检测到梯度中的NaN。

将CLIP在FP16和BF16中进行比较，可以发现BF16中梯度下溢仅偶尔发生。

在PyTorch 1.12及更高版本中，启用BF16只需对AMP进行小幅修改。

如果需要更多精度

另一种选择是TensorFloat32 (TF32) 数据类型。由Nvidia在Ampere GPU中引入，TF32是一种19位浮点数，结合了BF16的额外范围位，同时保留了FP16的精度。与FP16和BF16不同，TF32被设计为直接替代FP32，而不是在混合精度中启用。要在PyTorch中启用TF32，只需在训练开始时添加两行代码。

需要注意的是，在PyTorch 1.11之前，TF32在支持该数据类型的GPU上默认启用。自PyTorch 1.11起，为了保持跨GPU代的预期行为，必须手动启用。TF32的训练速度比BF16和FP16慢，理论FLOPS减半，但仍显著快于FP32训练。

BF16和TF32 可用于P4d、P4de、G5、Trn1和DL1实例。

解决问题于未然

上述示例展示了如何识别并解决FP16范围内的限制问题。然而，这些问题通常在训练后期才会出现。训练初期，模型产生较高的损失值，因此对异常值的敏感性较低。正如OpenCLIP训练中的情况一样，问题可能需要几天时间才会显现，这会浪费昂贵的计算资源。FP16和BF16各有优缺点。

FP16的范围限制可能会导致不稳定性并阻碍训练，而BF16精度较低，可能导致收敛性较差。理想情况下，我们希望在训练早期就能识别出易受FP16不稳定性影响的模型，以便在问题发生前做出明智的决策。

通过再次比较出现和未出现后续训练不稳定性的模型，我们可以识别出两个趋势：

• 使用FP16训练的YOLO模型和使用BF16训练的CLIP模型的梯度下溢率通常低于1%，并且随时间变化较为稳定。

• 使用FP16训练的CLIP模型在训练的前1000步内梯度下溢率为5%-10%，且随时间呈上升趋势。

我们发现，通过使用Tensor Collection Hook 跟踪梯度下溢率，可以在训练的前4-6小时内识别出梯度不稳定性上升的趋势。当观察到这种趋势时，我们建议切换到BF16。

结论

混合精度训练是训练最先进的大型基础模型的重要组成部分，但需要额外关注数值稳定性。尽管传统模型训练主要监控损失值和输出结果，我们的实验表明，了解模型的内部状态对于诊断混合精度数据类型极限问题至关重要。通过为模型添加Tensor Collection Hook，我们可以跟踪参数或梯度是否接近数值极限，并在不稳定性发生时或发生之前实施训练调整，从而可能节省数天的不成功训练时间。

要开始使用我们的Tensor Collection Hook，或复现我们的实验，请参阅我们修改后的OpenCLIP实现。

原文链接