关于人工智能:使用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

后果展现

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

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