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

后果展现

试验结果表明，这个办法在宽泛的模型中是无效的，并且能够开箱即用，不须要额定的超参数调优。

1、小规模的试验

第一组试验验证了在不同模型架构上的宽泛适用性。在FP32和FP16中训练了大量具备和不具备Unit Scaling的小型字符级语言模型，并比拟了后果。

在简直所有状况下，它都与基线性能匹配，甚至略有进步。当从FP32切换到FP16时，不须要调优。

2、大规模的试验

第二组试验在一个更大、更事实的生产级模型BERT[4]上验证了有效性。对单Unit Scaling模型进行调整，使其与规范BERT实现保持一致，而后应用来自英文维基百科文章的文本对其进行训练。

咱们对SQuAD v1.0和SQuAD v2.0评估工作的后果如下:

Unit Scaling可能取得与规范(基线)模型雷同的性能，并且在所有状况下都能够间接应用。基线模型和Unit Scaling模型并不完全相同，然而它们上游性能的偏差很小(Unit Scaling的BERT Base略低于基线，BERT Large略高于基线)。

FP8的实现是基于Graphcore、AMD和Qualcomm最近提出的标准化格局。Graphcore钻研之前证实了在FP8中训练损失缩放BERT而没有进化[5]，论文也证实了通过Unit Scaling也能够实现同样的成果。

要使FP8优于FP16，不须要额定的技术。只是简略地将matmul输出量化到FP8中，并可能精确地训练(FP8 E4变体中的权重和激活，以及E5中的梯度)。

低精度训练的将来

随着反对FP8的硬件在人工智能社区的采纳越来越多，无效、间接且有准则的模型缩放办法也变得越来越重要。Unit Scaling能够实用于宽泛的模型和优化器，并且计算开销最小。

下一代大型模型可能会宽泛应用低精度格局，所以这种缩放的办法十分的必要。低精度训练的效率劣势是微小的，Unit Scaling也证实了低精度并不一定会 升高模型的体现。

论文地址：https://avoid.overfit.cn/post/dfcaa9c45d70421a98f4df52a9e83610

援用

[1] P. Micikevicius et al., Mixed precision training (2018). 6th International Conference on Learning Representations

[2] O. Kuchaiev et al., Mixed-precision training for nlp and speech recognition with openseq2seq (2018), arXiv preprint arXiv:1805.10387

[3] P. Micikevicius et al., FP8 formats for deep learning (2022). arXiv preprint arXiv:2209.05433

[4] J. Devlin et al., BERT: Pre-training of deep bidirectional transformers for language understanding (2019). NAACL-HLT

[5] B. Noune et al., 8-bit numerical formats for deep neural networks (2019). arXiv preprint arXiv:2206.02915

本文作者：Charlie Blake