摘要
量化已成为降低大型语言模型 (LLM) 推理的计算和存储成本的主要方法。当前的大多数研究都集中在对权重和激活进行量化,以实现低比特通用矩阵乘法 (GEMM) 操作,而其余的非线性操作则以更高的精度执行。在我们的研究中,我们发现,在应用这些技术之后,LLM 推理中的主要瓶颈在于 softmax 层。softmax 操作包含三个阶段:指数计算、累加和归一化。我们的工作重点是优化前两个阶段。我们提出了一种分析方法来确定 softmax 函数输入的最佳剪切值,从而使 LLM 推理能够实现低于 4 比特的量化。这种方法可以加速 $e^x$ 和 $\sum(e^x)$ 的计算,同时精度下降最小或没有下降。例如,在 LLaMA1-30B 中,我们使用 2 比特量化在著名的“物理交互:问答”(PIQA)数据集评估中实现了基线性能。这种超低比特量化首次使累加阶段的加速达到约 4 倍。加速 $e^x$ 和 $\sum(e^x)$ 的组合导致 softmax 操作加速了 36.9%。