利用自动损失缩放更稳定地训练大模型

在本文中，我们将介绍一项由Graphcore拟未开发的原创技术，该技术可以在大模型的混合精度训练中，轻松而可靠地提高稳定性。它源于我们为IPU开发应用的独特经验，为了便于使用，现已被整合到Poplar SDK中。使在IPU上开发模型和进行实验的过程尽可能简单，以满足不同需求和不同熟悉程度的用户（从机器学习研究人员到数据科学家、MLOps工程师、云应用开发人员等），是拟未一贯的目标。

首先，我们将解释为什么损失缩放在大模型混合精度训练中的作用至关重要，尤其是考虑到推动机器学习进步的趋势。我们将展示此前损失缩放方法（手动和自动）存在的效率低下或容易失败的问题。然后，我们将展示拟未ALS算法是如何兼顾效率、易用性，以及高稳定性的，而这正是人们迫切需要的。

机器学习中的许多突破都是通过不断增加模型架构的大小和复杂性实现的。低精度的数值格式是克服伴随着模型大型化趋势而来的计算挑战的重要工具，它实现了几个好处：更低的存储占用、更大的带宽和吞吐量，以及由于降低了功耗而更节能的训练。然而，这些好处是以在训练过程中引入数值不稳定和降低一些模型的统计性能为代价的。对于目前的大型深度学习模型，从业者在训练过程中通常采用IEEE 32位和16位浮点的混合表示，这被称为混合精度训练。在本文中，我们将解释为什么需要损失缩放来实现混合精度训练的良好收敛，以及我们如何在训练过程中通过观察梯度直方图的方式来自动调整它。

将精度从IEEE float-32降低到IEEE float-16，缩小了激活、权重和梯度的动态范围。确保它们的信号不会因为float-16的下溢而丢失是至关重要的。同样地，模型可以容忍一定程度的溢出，但如果溢出过多，就会导致不稳定和突然故障。总而言之，下溢和溢出都会阻碍训练收敛。谷歌推出的bfloat-16是float-16精度的一个替代方案，通过减少尾数位数来保持float-32的动态范围。bfloat-16减少了梯度下溢或溢出的风险，而且不需要损失缩放，但这是以可能妨碍统计收敛为代价的。这在谷歌的Gopher训练论文^[1]中有所体现。该论文的作者发现，与float-32相比，即使是用随机舍入等技术来补充，使用bfloat-16仍然会降低性能。

损失缩放的目的是在整个动态范围内转移梯度分布，在float-16中（尽可能地）防止下溢和溢出。顾名思义，梯度的缩放是通过将损失函数乘以一个常数系数来实现的。因此，通过反向传播获得的梯度也是按这个常数缩放的。在权重更新过程中必须考虑到这种缩放，否则训练会受到损失缩放值的影响。下图总结了SGD优化器的损失缩放方式。

使用SGD优化器的损失缩放

众所周知，大模型在训练过程中会出现不稳定性，这可能是数值问题甚至是硬件故障造成的。然而，在科学出版物中通常不强调这个实际的话题：人们把主要的兴趣点放在了模型的验证分数上，却很少有人关注多次尝试以稳定的方式训练大模型所涉及的挑战。

在混合精度训练中调控损失缩放正是这种挑战的一个典型例子。调整可以手动完成，也可以通过动态或自动程序完成，该程序根据一定的标准在每一步骤或每个步骤间隔内选择一个损失缩放值。文献中提出了几种设计自动程序的方法。一种方法是基于计算中出现的NaN值，最近Meta的工程师在OPT-175B模型中采用了这种方法。然而，快速回顾一下Meta训练日志^[2]中提到的损失缩放就会发现，在许多尝试中，由于损失缩放意外激增，导致损失函数突然不稳定，缺乏收敛性，无法进行稳定的训练。他们的动态损失缩放程序的目的其实是为了稳定训练，而这种现象是该程序的意外行为。

拟未在理解大模型通常表现出的故障模式方面很有经验，并已经完善了这种损失缩放程序。这使我们能够通过在Poplar SDK中开发高效的内部工具来加强大模型的训练稳定性。在讨论拟未ALS算法的细节之前，让我们首先看看BERT Large是如何由于错误地选择静态损失缩放而失败的。

论文《面向深度学习的大批优化：在76分钟内训练BERT》^[3]中介绍了LAMB优化器，我们按照这篇论文中的超参数在IPU-POD₆₄上预训练BERT Large的阶段1。我们对所有的权重、激活和梯度使用float-16表示。只有LAMB优化器的第一和第二时刻被设置为float-32。您可以在我们的GitHub存储库^[4]中找到我们基于Hugging Face transformers的开源实现以及关于数据集的细节。

下图描述了配有四个恒定的不同损失缩放值的BERT Large的训练演变。对于1、2⁷和2²²的损失缩放值，无论选择什么种子，训练都会出现发散。对于2⁷和2²²之间的值，运行收敛取决于种子，收敛的概率较低或较高。特别是将损失缩放设置为2¹⁵将带来最高的收敛率为90%（即只有10%的种子在训练中最终出现发散）。为了简单起见，我们只绘制了一个发散和收敛的种子，但对许多其他种子也重复了实验。这意味着，在使用静态损失缩放时，无论选择什么损失缩放值，都可能存在发散的风险。然而，选择正确的损失缩放值对于确保尽可能降低发散率是至关重要的。

我们观察到，损失缩放等于或低于1，以及高于或等于2²²的运行总是很快就失败了。这是由于在float-16中出现了明显的梯度下溢/溢出，并且运行无法在训练开始之后继续。损失缩放等于2⁷的运行在2k个训练步骤之前进展充分，这时学习率调度器达到了峰值（如下图），运行变得更加不稳定。这时候，运行在短短的几个训练步骤中就灾难性地失败了。这也正是带有损失缩放2¹⁵的运行在10%的种子中失败的时候。像这样的突然失败是大模型的特点，社区里有一些讨论正是关于这些失败是如何造成的，是否涉及梯度下溢或溢出，以及哪些与权重或激活有关的梯度与其相关。

BERT Large阶段1预训练的MLM精度，适用于各种损失缩放选择

BERT Large阶段1预训练的学习率调度器

很明显，在任何情况下，对于训练稳定性来说，找到适当的损失缩放都至关重要：所有显示的BERT运行都共享相同的超参数，但在损失缩放系数和种子方面有所不同。优化损失缩放可以带来大多数种子的收敛。

我们还应该注意到，找到适当的损失缩放值的代价可能相当昂贵，要消耗几个小时甚至几天的计算时间，具体时间长短取决于所使用的系统。如果需要对多次损失缩放范围重复此操作，那么消耗的计算资源就会大大增加。

这个问题可以由下文所述的拟未ALS算法解决，它可以使我们在上文所示的BERT训练配置中实现100%的收敛率。

拟未ALS方法基于从权重和激活方面对梯度的观察。这些观察结果被用来生成直方图，为调整损失缩放系数提供信息，目的是防止梯度分布中的过度下溢或溢出。这种方法与其他策略不同，其他策略通常是根据溢出事件发生时计算中出现的NaN/Inf值来调整损失缩放系数。我们认为观察梯度是一种更明智的策略，因为它可以平衡溢出和下溢的数量。

我们将梯度聚集到直方图中，只有两个仓——h1和h2，它们涵盖了float-16的全部动态范围（即从0到65504）。每当有溢出事件发生时，数值就会在最大的数字表示上被剪裁，在float-16中对应的是65504。由于这种剪裁方式，如果我们只关注最终值，可能会错过中间操作中发生的溢出事件。为了同时考虑到这些中间的溢出，我们将两个仓之间的仓边缘设置为接近溢出但留有一定的余地，在实践中，像2¹³这样的仓边缘已经带来了稳定的运行。我们在每次权重更新时计算两个仓之间的比率，根据10^-7的阈值，我们将损失缩放值增加一倍或减少一半。下图直观地描述了我们的ALS算法。

ALS算法图解

该算法背后的逻辑很简单：损失缩放不断翻倍以防止下溢。只有当上仓计数与两仓之和相比达到一定比例时，损失缩放才减半以避免过度溢出。

可以调整该算法，以减少其计算开销：可以引入一个周期，这样做是有意义的，因为对于大多数模型来说，梯度分布变化缓慢，意味着不需要在每个训练步骤后调整损失缩放系数。

另一个关系到追踪梯度的可能修改：仅选取某些层/操作的梯度、权重梯度、激活梯度等，而非所有的梯度，这可能会导致ALS方案依赖模型，可能使其无法普遍适用。为简化说明，在这篇文章中，我们将周期设置为1，并跟踪所有梯度。

为了测试拟未ALS算法的易用性和稳健性，我们将其应用于近年来颇为流行的两个深度学习模型的预训练中：用于语言的BERT Large（BERT）和用于视觉的EfficientNet-B4。

BERT（如前面的静态损失缩放的例子）是根据论文《面向深度学习的大批优化：在76分钟内训练BERT》^[5]训练的。EfficientNet-B4的基础遵循了参考论文《EfficientNet：重新构想面向卷积神经网络的模型缩放》^[6]，不过采用了《让EfficientNet更高效：探索与批无关的归一化、分组卷积和降低分辨率的训练》^[7]一文中提出的算法改进。

在这两个模型中，除了优化器的状态是float-32，我们将所有的权重和梯度保持在float-16。部分操作也是以float-16进行的。BERT的优化器是LAMB，EfficientNet的优化器是RMSProp，学习率遵循初始阶段预热，直到达到最大值（BERT第一阶段为2k个训练步骤，BERT第二阶段为275个步骤，EfficientNet为4个训练epoch），随后BERT的学习率呈线性衰减（EfficientNet为指数衰减），直到训练结束。

您可以在我们的GitHub存储库^[8]中找到拟未基于Hugging Face Transformer的BERT开源实现，以及关于维基百科数据集的细节。同样，我们内部开源的EfficientNet实现，以及ImageNet数据集的细节，也都位于我们的GitHub存储库^[9]中。

BERT的训练图直接显示在下面的图表中。训练BERT包括两个阶段：第一个阶段使用序列长度128，第二个阶段使用512。这主要是为了提高效率：注意力在序列长度上是二次的；用较短的序列长度训练比较便宜，但学习位置嵌入和长距离依赖需要用较长的序列长度训练。

我们观察到，在第二阶段结束时，损失函数达到了1.24的值。这是一个SOTA得分，使我们在SQuAD中进行微调时达到84.38的精确匹配和90.92的F1（BERT论文^[10]的参考值分别为84.1和90.9）。关于损失缩放图，我们最初将损失比缩放值设置为1，但接下来在每一个训练步骤之后都会进行调整，从而导致了一个之字形的轮廓，在第一阶段的215=32,768和第二阶段的214=16384之间震荡。因为损失缩放更新周期被设置为1，这个之字形的轮廓形成于每个训练步骤后，但我们也可以使用频率更低的更新。

有趣的是，我们看到对于第一阶段，除了在1k和2k之间的一些步骤中（这是学习率在之前使用静态损失缩放的预训练示例中达到最高值的地方），损失缩放值在整个训练过程中保持振荡和某种程度的恒定。拟未ALS算法对梯度分布的变化做出了反应，这可能是由于分布中的变异较大或总体值较大。在任何情况下，该算法都会降低损失缩放值，以防止过度溢出。

使用ALS的BERT Large预训练阶段1的损失演变

使用ALS的BERT Large预训练阶段1的损失缩放演变

使用ALS的BERT Large预训练阶段2的损失演变

使用ALS的BERT Large预训练阶段2的损失缩放演变

EfficientNet的训练图如下图所示。我们没有绘制每个训练步骤的图，而只是绘制了每个训练epoch的损失和损失缩放的最终值。损失函数在训练过程中平稳下降，最终的验证精度为82.33，与《让EfficientNet更高效：探索与批无关的归一化、分组卷积和降低分辨率的训练》^[11]一文中的参考值82.3一致。

损失的最终下降是由于将最后一个检查点作为之前所有检查点的指数移动平均值加以计算。关于损失缩放演化，恒定的形态是因为每个训练epoch只绘制一次；实际上，损失缩放在每次权重更新后仍在更新，其方式与上述BERT类似。在整个训练过程中，最优损失缩放值也存在一些变化。

使用ALS的EfficientNet-B4预训练的损失演变

使用ALS的EfficientNet-B4预训练的损失缩放演变

虽然损失缩放对于在动态范围内适当转移梯度和防止下溢至关重要，但根据模型和配置不同，这可能还不够。为了提高稳定性，拟未使用随机舍入、累积的运行平均值和某些元素（如优化器状态时刻）的float-32精度来补充损失缩放。

IPU原生支持的随机舍入可用于混合精度训练，以帮助缓解使用float-16部分时的精度损失（在16.16 AMP中），或在没有主权重的float-32副本的情况下进行训练。在随机舍入中，产生哪种输出是不确定的。例如，当随机舍入1.2时，20%的时间里结果是2，80%的时间里结果是1。在数学上，对于一个量化步长为Δ的输入x，随机舍入公式为：

随机舍入的好处是，尽管增加了一个可容忍的误差水平，但产生一个无偏的量化（其中$\mathbb{E}\{ SR (x)\}=x$）。这意味着，在许多这样的添加中，被添加的量化噪声的平均值为零，而注入的携带数接近于在更高精度下通过累积传播的情况。

作为一个例子，让我们来看看在ResNet32中使用CIFAR-100、SGD和批尺寸为4的随机舍入的影响。下图比较了不同的float-32和float-16精度的组合对算术和权重更新格式的验证精度。启用随机舍入可以防止在算术和权重更新格式中采用float-16时验证精度的下降。

为计算梯度的平均值，数据并行性和梯度累积都需要进行累加操作。如果在计算平均值时首先将所有梯度相加，然后再除以梯度总数，那么在float-16中就有溢出的风险。

因此，我们实现了一个运行平均值累积，确保累积的梯度不会大于数据并行性和梯度累积的最大微批梯度的大小。

具体而言，给定梯度$g_{1},\dots, g_{N}$，以$M_0=0$通过\[M_k = \frac{k-1}{k} M_{k-1} + \frac{1}{k} g_k\]迭代计算出最前的$k$个梯度的平均值。

事实上，运行平均数对于将损失缩放的调整与累积规模脱钩来说是必要的，因为有了它，梯度和它们的累积都有相同的数量级。