使用Unit Scaling进行FP16 和 FP8 训练

Unit Scaling 是一种新的低精度机器学习方法，能够在没有损失缩放的情况下训练 FP16 和 FP8 中的语言模型。

使用FP16和BFLOAT16替代FP32可以将内存、带宽和计算需求的大幅减少，这也是目前越来越大的模型所需要的。

背景介绍

随着支持fp8的硬件的发展，在不影响效率的前提下，进一步降低精度也成为了可能。但是这些较小的、低精度的格式在实践中并不总是易于使用。对于FP8来说则更加困难。因为这些较小的格式通常将用户限制在更窄的可表示值范围内。为了解决这个问题，Graphcore Research开发了一种新方法，我们称之为Unit Scaling。

上图为FP16和FP8中量化的不同尺度的正态分布的信噪比(SNR)。对于较小的数字格式，信号在较窄的尺度范围内较强。

Unit Scaling是一种模型设计技术，它在初始化时根据缩放原则进行操作：也就是说对激活、权重和梯度的单位方差进行缩放。模型会自动生成针对低精度数字格式进行良好缩放的张量。并且使用更简单，并最大限度地减少这些表示的缺点，与低精度训练的替代方法不同，它引入的开销和额外的复杂性很小。

论文的方法取得了突破性的成果：首次在 FP16 甚至 FP8 中准确地训练了 BERT Base 和 BERT Large 模型，并且没有缩放的性能损失。模型也不需要额外的超参数，可以直接使用。

对于关心结果并因此希望在 FP16 和 FP8 中进行训练的人来说，Unit Scaling提供了一个直接的解决方案。

FP16/FP8训练的现有方法

FP16和FP8训练需要某种形式的缩放来保持值在范围内。目前的做法如下:

1、(静态)损失缩放

缩小范围对于训练期间的向反向传播是具有挑战性的通常会导致梯度下溢。为了解决这个问题，最常见的方法是将损失乘以超参数以增加梯度的大小 [1]。由于没有原则性的方法来提前选择损失的规模，所以这个超参数通常需要多次运行。

2、自动损失缩放

通过基于运行时的梯度溢出（或直方图）[2] 动态调整损失比例，可以避免超参数扫描的需要。但是这种自动方案会增加开销和复杂性。

3、张量缩放

上述方法的另一个缺点是它们只提供单一的全局损失尺度。另外一种解决方案是根据张量统计 [3] 重新缩放值。这也是一种自动/运行时方案，很复杂且难以有效实施。

Unit Scaling

Unit Scaling 在前向和反向传播中引入局部缩放因子控制值的范围。选择的范围是基于每个操作符如何影响值规模的理解，而并不是使用运行时分析得到的。通过选择正确的比例因子，每个操作都大致保持其输入的比例。通过将其应用于所有操作，可以控制整个模型中传播初始（单位）比例，从而实现全局的缩放。

这种方法比自动缩放方案更简单，因为唯一的额外开销是应用缩放因子。对于 BERT Large，这会将 FLOPs 增加 0.2%，应该可以忽略不计。

模型可以通过应用以下方法进行Unit Scaling:

用单位方差初始化无偏差参数
计算所有操作的理想比例因子
识别非切边并限制使用它们的操作具有相同的缩放比例
用加权的加法替换加法

下面我们将更详细地解释这些规则。

1、理想的比例因子

我们可以对一些操作进行数学分析，以确定它们如何影响输入的方差。

比如基本矩阵乘法 XW（其中 X 是 (b × m) 矩阵，W 是 (m × n) 矩阵）的输出方差为 σ(X)² · σ(W)² · m。要缩放此操作，我们必须确保 σ(X)² = σ(W)² = 1，然后将 1/√m 乘法添加到输出。

对于反向传播，需要引入了两个新的矩阵乘法，理想的比例因子为 1/√n 和 1/√b。其他操作也可以类似分析，输出方差不容易分析，所以可以使用经验方法来找到缩放因子。

在论文作者中提供了更详细的分析，以及常见操作的概要及其理想的比例因子。

2、切边

直接将这些理想的比例因子应用于正向和反向传播中会产生无效的梯度。为了避免这种情况，某些操作需要使用共享的缩放因子。

我们使用前向计算图并找到所有没有用切边表示的变量(如果去掉这些边，会将图分割成两个不相连的更小的图)。比如，下面是一个transformer的FFN层:

在权重、输入和输出变量上有切边。该图还显示了为第二个matmul的反向传播生成的梯度操作(我们只考虑正向图的切边)。

因为 x₃ 不是切边，所以可以限制 ∇x₃ 的 matmul 使用与前向传播中相同的比例因子，但是由于 w2 是切边，它允许有自己的反向缩放因子，所以为受约束的操作选择共享比例因子，采用之前计算的理想比例因子的几何平均值。

这个规则听起来很复杂，但实际上它通常可以归结为一个简单的过程:为权重梯度提供它们自己的比例因子（也就是模型中的任何编码器/解码器层）。

3、加权加法操作

最后一步是用加权的加法替换加法操作。根据设计的单位缩放产生的变量具有相等的尺度，如果我们将两个张量相加，它们实际上都具有相等的权重。但是在某些情况下，例如残差连接就需要一个不平衡的权重来获得良好的性能。所以将加法操作替换为加权(和单位缩放)加法等效操作。

对于残差连接，可以推导出以下方案:

代码实现

下面的代码展示了一个在PyTorch中实现Unit Scaling的FFN层。

首先定义创建基本操作的缩放版本，例如scaled_projection:

 classScaledGrad(autograd.Function):
   @staticmethod
   defforward(ctx, X, alpha, beta):
     ctx.save_for_backward(tensor(beta, dtype=X.dtype))
     returnalpha*X
 
   @staticmethod
   defbackward(ctx, grad_Y):
     beta, =ctx.saved_tensors
     returnbeta*grad_Y, None, None
 
 defscaled(X, alpha=1, beta=1):
   """forward: Y = X * alpha, backward: grad_X = grad_Y * beta"""
   returnScaledGrad.apply(X, alpha, beta)
 
 defscaled_projection(X, W):
   (b, _), (m, n) =X.shape, W.shape
   alpha=beta_X= (m*n) **-(1/4) beta_W=b**-(1/2)
   X=scaled(X, beta=beta_X)
   W=scaled(W, beta=beta_W)
   returnscaled(matmul(X, W), alpha)

这样我们就可以创建完整的层。我们只演示一个标准FFN和它的缩放版本:

 classFFN(nn.Module):
   def__init__(self, d, h):
     super().__init__()
     self.norm=LayerNorm(d)
     sigma= (d*h) **-(1/4)
     self.W_1=Parameter(randn(d, h) *sigma)
     self.W_2=Parameter(randn(h, d) *sigma)
 
   defforward(self, X):
     Z=self.norm(X)
     Z=matmul(Z, self.W_1) Z=gelu(Z)
     Z=matmul(Z, self.W_2) returnX+Z
 
 
 classScaledFFN(nn.Module):
   def__init__(self, d, h, tau):
     super().__init__()
     self.norm=ScaledLayerNorm(d)  # Not defined here
     self.W1=Parameter(randn(d, h))
     self.W2=Parameter(randn(h, d))
     self.tau=tau
 
   defforward(self, X):
     a= (1-self.tau) ** (1/2)
     b=self.tau** (1/2)
     Z=self.norm(scaled(X, beta=b))
     Z=scaled_projection(Z, self.W1)
     Z=scaled_gelu(Z)  # Not defined here
     Z=scaled_projection(Z, self.W2)
     returnX*a+scaled(Z, b)  # fixed(𝜏) weighted add