总结
LLM 4 位量化是在设备上部署的一项重要技术,因为它可以大大减少推理过程中的内存使用量。然而,众所周知,训练过程中出现的 "异常值"(极端激活值)会大大降低量化的性能。传统的方法是在训练后处理这个问题,例如通过训练后量化(PTQ),但这些方法并没有从根本上解决这个问题,而是将异常值视为不可避免的。
在本研究中,我们提出了解决这一问题的新视角:离群值安全预训练(OSP)训练框架,它由 Muong 优化方法、单尺度 RMS 归一化(SSNORM)和可学习的嵌入式投影(EMBPROJ)组成。与传统模型相比,OSP 通过在万亿令牌规模上训练 1.4B 参数模型,显著提高了量化容忍度,减少了性能下降。
建议的方法
所提出的离群值安全预训练(OSP)是一种预训练方法,旨在从根本上防止离群值。该框架由三个部分组成。
首先,Muon 优化方法与 Adam 和 AdaFactor 等传统的对角预条件优化方法不同,它使用了一种正交梯度矩阵的算法。这可以防止激活集中于特定通道(特权基础),并实现所有通道的平等学习。
其次,Single-Scale RMSNorm(SSNORM)消除了传统的按通道缩放,对所有维度使用单一缩放因子,防止了归一化带来的偏差。这使他们能够减少训练过程中的不稳定性,同时也避免了对激活的过度抑制。
第三,嵌入投影(EMBPROJ)使用可学习的投影矩阵均匀化激活分布,以防止嵌入层产生局部异常值;只有嵌入层使用 Adam 进行训练,其他参数使用 Muon。此外,为了实现实用性和计算效率,我们还采用了一种分离式优化策略,即只使用 Adam 对嵌入层进行训练,而对其他参数则使用 Muon。
实验
实验是在一个 14 亿参数的 LLM 上进行的,使用 1,000 亿和 1 万亿代币进行可扩展训练。首先,利用峰度过大量化异常值,并观察其演变情况。使用传统 Adam 训练的模型显示出峰度超过 1000 的离群值,而使用 OSP 的模型则继续保持 0.04 的极低值。
为了验证 4 位量化的性能,比较了 10 个基准(如 ARC、MMLU、GSM8k 等)的平均得分;Adam 训练的模型得分 26.5,而 OSP 模型得分 35.7。此外,当与训练后量化(PTQ)方法结合使用时,OSP 模型始终保持较低的复杂度,并显示出强大的量化性能。
此外,我们还分析了注意力机制中的 "注意力下沉 "现象,这种现象在异常值消失后依然存在,表明两者是由不同的机制造成的。这证明了 OSP 作为一种优化的量化训练方法的有效性。