随着大型语言模型（LLMs）参数的不断增加，数百亿甚至数千亿参数，对专用硬件加速器制造商产生了巨大压力，使后者的创新设计成为人工智能行业中发展最快的领域之一。
为了在计算和存储有限的情况下有效且准确地处理LLMs，已经探索了各种方法。在这些方法中，各种量化技术已成为社区的主要焦点，作为减少计算、通信和存储需求的一种手段。降低精度自然带来许多挑战，因为可用值表示的范围有限。
在处理硬件上的流行Transformer模型时，一个主要问题是计算LayerNorm，因为累积方差需要比硬件允许的更宽的动态范围。在这篇文章中，作者解决了这个问题，并提出了一种计算高效的扩展技术，可以轻松应用于推理过程中的Transformer模型。
作者的方法提出了一种基于立即前线性层静态权重的LayerNorm输入的简单扩展方法。
扩展因子在离线计算，仅基于线性层权重，因此推理过程中没有延迟或计算开销。
最重要的是，作者的技术确保在计算过程中不会发生数值问题，如溢出或下溢。这种方法提供了一种在各种硬件架构上实现平稳、准确和资源有效的推理方法。本文还提供了理论证明和数值模拟的支持。
1 Introduction
最近，基于Transformer的预训练语言模型（LLMs）已经成为主导深度神经网络（DNN）架构，原因在于其在所有语言建模[2; 3]、文本处理[4]、图像和视频生成[5]等任务中表现出了前所未有的性能。然而，这种成功是以巨大的计算、存储和数据传输为代价的。在过去的几年里，为了满足LLM训练和推理的需求，专门用于加速硬件的行业已经崭露头角[6; 7]。另一个旨在使推理可行和可持续的主要举措是开发低精度格式[8; 9; 10]，高效的量化技术[11]，算法解决方案[12]，精确的近似[13]，以及其他软件优化[14; 15]。
高效量化技术如GPTQ [16], AWQ [17], SmoothQuant [18], KVQuant [19], K-sort [20],以及许多其他技术，使得可以使用低精度格式存储和处理LLMs。通常，这涉及在FP32格式下训练模型，然后在部署到推理硬件之前将其转换为4、8或16位精度格式[11; 21; 20]。最流行的方法是将静态权重压缩为4或8位整数或浮点数，并将激活值减少为FP16或BF16[22]。在本论文中，作者专注于在FP16激活上运行的加速器家族，因为它们在当前Transformer中的广泛使用 [23; 24]，并特别针对FP16相对较窄的动态范围（可表示数的范围）可能带来的显著计算挑战。这个问题在LayerNorm计算过程中最为关键。重要的是，由于它们在训练过程中防止梯度爆炸或衰减，因此在当前Transformer中包含数十个甚至数百个LayerNorm运算符是无法避免的。然而，在推理阶段，处理LayerNorms在加速器上的计算是非常具有挑战性的，因为它们需要对输入进行累积平方以进行方差（和范数）计算 [26]。在FP16中累积如此大量正值几乎肯定会导致溢出。
在这项工作中，作者解决了这个问题，并提出了一种高效、理论上合理且易于实现的可扩展技术，该技术可导致LayerNorms中的FP16溢出（或下溢）问题的完全消除。首先，请注意，LayerNorm输入的缩放不会影响输出，因为归一化操作的均匀性，但可以非常显著地改变分母计算中的累积数值范围。基于这个观察，作者开发了SLaNC（静态LayerNorm校准）方法，它为任何Transformer的所有LayerNorm的输入提供了简洁的封闭公式。重要的是，SLaNC计算的缩放仅基于前一层静态权重的线性层，因此可以在不影响推理时间的情况下离线计算。SLaNC提供的公式理论上由推导和详细解释进行证明，仅涉及可以直接且精确计算的静态权重矩阵的范数。
本文其余部分的组织如下。首先，作者给出了符号说明。在第二部分，作者提出了由于FP16中的LayerNorm计算而产生的数值问题。第三部分介绍了SLaNC技术及其理论依据。在第四部分，作者通过LLM的Llama家族进行了数值模拟。第五部分是结论部分。
注解 在本文中，以下注解被使用。矩阵用大写粗体字母表示，向量用小写粗体表示。矩阵和的适当尺寸的乘积用或表示，而它们按元素乘积则用表示。对于矩阵,作者用表示其Frobenius范数，用表示其谱范数；对于向量,用表示其欧几里得范数。给定向量,作者用表示具有主要对角元素的对角矩阵。
2 Problem Formulation
将一个LLM量化到低精度格式（如，4位、8位或16位）会导致输出质量显著降低，因此必须与能够恢复准确性的先进技术相结合应用[16; 17; 18; 19; 20; 27; 28]。然而，将模型导入低精度格式面临更大的挑战，即这些格式的动态范围有限，如果盲目应用，可能会完全破坏计算流程。最典型的例子是计算LayerNorm，由于不可避免的溢出和下溢，在FP16加速器上变得不可能。
LayerNorm Compute
Layer Normalization（LayerNorm）已成为现代深度神经网络（DNNs）中最常见的非线性操作之一，因为它在训练过程中防止了梯度衰减或爆炸。大量文献已经证明，当前的DNN架构在不需要频繁对隐藏状态进行规范化的情况下是无法实际训练的[29; 30; 31]。最新的Transformer模型包括数十个甚至数百个LayerNorm操作符，它们被引入以方便训练，但由于计算它们的分母产生的数值问题，使得推理变得困难。
根据公式2，标准计算σ的方式需要将输入向量元素的平方和。这些元素的范围决定了这种累加是否容易导致FP16或FP8格式下的溢出或下溢。值得注意的是，大多数可用的LLM加速器在非线性模块中仅以FP16格式处理非线性操作[32, 33]。虽然一些加速器在非线性模块中支持FP32累加，但这种选项通常会导致高延迟增加，使得FP32区域变得不切实际。
图2(a)和图2(c)显示了Llama-2的一个层中公式2中平方和的典型分布。作者观察到在许多情况下，计算结果的值超出了FP16的范围，导致无效的推理。
请注意，Transformer架构有两种味道，取决于残差分支的位置。它可以位于LayerNorm（预LN残差）之前（pre-LN残差）或之后（post-LN残差），如图1所示。最初，建议使用post-LN选项[1]，但后来另一种由于观察到训练速度更快而变得非常流行[34]。具体而言，由于缺乏空间，作者在这里主要关注post-LN设计，但作者强调推导和结论同样适用于pre-LN。
3 LayerNorm Scaling
Dynamic Model Calibration
自然解决LAYER NORM计算中的溢出或下溢问题的方法是适当缩放其输入。确定正确的缩放因子似乎具有挑战性，因为在避免溢出的同时，作者也不想过度丢弃输入导致下溢，反之亦然。因此，任何合理的缩放算法都必须考虑实际LAYER NORM输入值，而不能盲目地设置缩放参数。
一种常见的解决方案是校准缩放因子。这涉及将一个测试数据集通过Transformer，以评估输入向量范数的范围，并根据这个范围的一些一阶统计量（例如均值或中范数）设置缩放因子。这种技术需要额外的校准数据和显著的计算开销，即使对于像LayerNorm这样的基本操作也是如此，因此这种方法实用性不强。
Analytical Static Scaling
在这项工作中，作者提出了一种不同的方法，使得可以对所需的缩放因子进行离线分析计算。缩放因子仅基于当前LayerNorm之前的前线性层静态权重来确定。这样，作者可以静态地校准一个模型中的所有LayerNorm运算符，而无需使用校准数据集或额外的运行时资源 -- 所有计算都在模型编译期间预先完成。
该方法的思想基于一个简单的观察，即Transformer中的LayerNorm经常且按规律出现，因为任何大型Transformer都是一个由多个相同解码器组成的链。通常，连续的两个LayerNorm围绕着每个解码器的注意力或MLP（多层感知机）块。等式1表明，作者可以将LayerNorm视为一个欧几里得归一化 followed by a diagonal matrix multiplication。1 从LayerNorm操作的自然分解中，作者可以推理出，在归一化（LayerNorm的第一步）之后，隐向量的范数等于1。作者的目标是追踪从这一点到下一个LayerNorm的计算图，并衡量在经过过程中的变化，从而评估其数量级。
SLaNC for Standard MLP Block
请提供需要翻译的英文AI学术论文内容。
SLaNC for Llama MLP Block
采用相同的算法，作者得到了Llama家族模型解码器的改进MLP块所对应的L LayerNorm的缩放因子的类似公式，如图1(b)所示。在这里，除了标准的MLP块的两个线性层之外，还有一个线性层，其输出与元素乘法的方式非线性层的输出相乘。非线性函数本身通常选择GeLU。后MLP块LayerNorm的输入y的读法如下：
类似上述原理，结合矩阵范数的基性质，可以得到
SLaNC for the Attention Block
接下来，作者推导出LayerNorm的缩放因子公式，该公式遵循具有h个头的标准注意力块。
并且得出在后自注意力 LayerNorm 操作中应使用以下缩放系数。
为了展示作者SLaNC扩展技术的威力，作者提出了Llama模型的模拟结果。请注意，Llama架构替换了LayerNorm，采用了根均方层归一化（RMSNorm）[35]，这和前者的区别仅在于方程2中的均值μ减法被省略，因此不会影响SLaNC扩展。
在作者的第一次实验中，作者收集了没有缩放和具有SLaNC缩放的RMSNorm算子分母的平方和的实证统计数据。为此，作者将Llama-2-7b应用于Wikitext 2数据集。图2(a)和2(c)显示了该模型连续两个RMSNorm的特征典型的直方图。
作者发现，在许多情况下，平方和显著超过了FP16范围，导致溢出。SLaNC缩放则彻底改变了这种情况，不仅将直方图移入FP16范围，而且还保持了范围两端的 safe margins，如图2(b)和2(d)所示。
接下来，作者将Llama模型在相同的Wikitext 2数据集上与RMSNorm的默认FP32实现以及FP16（所有其他设置保持不变）的平方累加求和进行比较。表1显示，在FP16累加时，由于许多溢出，出现了显著的降级。这个问题在应用SLaNC缩放后完全得到解决。作者还注意到，在所有标准模型中，LayerNorm或RMSNorm操作的输入方差分母中增加了一个小常数ε。
这样作者可以在下标溢出时避免除以零，并提高数值稳定性。由于SLaNC缩放是在事先知道的，作者也可以轻松地将它们应用于ε常数（实际上，作者将ε除以平方的SLaNC缩放）。正如表1的底部一行所示，现在FP16 SLaNC缩放可以精确地再现默认FP32值。
结论
在本文中，作者提出了一种新颖的SLaNC技术，使得在FP16加速器上进行LLM推理成为可能，而无需将LayerNorm运算符转换为FP32。
这一基于理论的方法为LayerNorm输入的缩放因子的离线计算提供了易于使用的公式。
SLaNC缩放因子确保了在FP16中精确计算LayerNorm，并且可以证明避免溢出和下溢。
通过保持所有计算在FP16中，SLaNC算法实现了低延迟的精确计算，这一点通过作者的广泛数值模拟得到了证明。
参考文献
[0]. SLaNC: Static LayerNorm Calibration.
END
作者：小书童
来源：集智书童
推荐阅读
- LLM 量化新篇章，4-bit 权重激活量化几乎无损！FlatQuant 的平坦之道
- 【ml-engineering 翻译系列】大模型推理
- 【翻译】使用PyTorch FSDP和Torch.compile最大化训练吞吐量
- 超越SOTA ！ YOLOv8-ResCBAM 集成注意力机制以提高检测性能 ！
欢迎大家点赞留言，更多Arm技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
