跳到主要内容
起居室老虎
返回

I-LLM: Efficient Integer-Only Inference for Fully-Quantized Low-Bit Large Language Models

1,219 字约 5 分钟

原文:arXiv:2405.17849 · Hu et al., 2024

摘要: 后训练量化(PTQ)是一种加速大语言模型(LLM)推理的有效技术。然而,现有工作在推理过程中仍需要大量浮点(FP)操作,包括额外的量化和反量化,以及 RMSNorm 和 Softmax 等非线性操作。这一限制阻碍了 LLM 在边缘设备和云设备上的部署。在本文中,我们发现 LLM 实现全整数量化的主要障碍在于线性和非线性操作中跨通道和跨标记的激活大幅波动。为了解决这个问题,我们提出了 I-LLM,一种为 LLM 量身定制的全整数全量化 PTQ 新框架。具体来说,(1) 我们开发了全平滑块重建(FSBR),以积极平滑所有激活和权重的跨通道变化。(2) 为了缓解跨标记变化引起的性能下降,我们引入了一种称为动态整数矩阵乘法(DI-MatMul)的新方法。该方法通过动态量化输入和输出的全整数操作,实现了全整数矩阵乘法的动态量化。(3) 我们设计了 DI-ClippedSoftmax、DI-Exp 和 DI-Normalization,利用位移操作有效执行非线性运算,同时保持精度。实验结果表明,我们的 I-LLM 在精度上可与 FP 基线媲美,并优于非整数量化方法。例如,I-LLM 能够在 W4A4 设置下运行,且精度损失可忽略不计。据我们所知,我们是首个弥合全整数量化与 LLM 之间差距的研究。我们已将代码发布在 anonymous.4open.science,旨在推动该领域的发展。

1. Intro

之前的 LLM PTQ 经常要做伪量化,做实际计算的时候还是用的浮点型。虽然降低了数据传输和带宽的需求,但是实际上没有减少计算量(因为计算的时候还是 float)。

LLM 量化的 Integer-Only 的工作很少,之前的工作主要集中在 CNN,后来有对 Bert 和 ViT 的 Integer-Only 的量化。LLM integer only 的工作这是第一篇。

Contribution:

Untitled.png

LLMs Quantization: 从优化矩阵乘、activation outliers 到 KV Cache、ZeroQuant 这样包含一部分系统的工作,QAT 和 PTQ 都有。这篇文章是一个 PTQ 的工作,对于超大的 LLM 来讲 QAT 可能带来的 fine-tuning 开销太大了。

Integer-only Quantization: dyadic arithmetic on CNN and transformer,比如 I-BERT,但是在比如 GELU 这样的操作里面还是会包含 FP 计算。I-ViT 是 Integer-Only 的 ViT,但是在 LLM 的数据和情况上是否有效还不清楚。

3. Method

Challenges of Integer-Only Quantization for Large Language Models:

  1. Quantizing the activation of LLMs, especially those originating from non-linear operators, poses a formidable challenge. As evidenced in Fig 1, the divergence in these non-linear activations often surpasses that of linear operators, particularly pronounced in models such as LLaMA. Previous methods have failed to address these non-linear activations, and straightforwardly quantizing non-linear layers may lead to substantial accuracy degradation.
  2. prior integer-only quantization techniques have overlooked the distinctive traits of LLMs , including divergent activation scales and the substantial overhead of loading large-scale weights. Even the W8A8 method introduced in I-BERT can lead to catastrophic outcomes, as shown in Fig 4, let alone more aggressive quantization methods like W6A6 or W4A4.

Untitled.png

Untitled.png

3.1. Basic Mathematical Notations

3.2. Fully-Smooth Block-Recontruction(FSBR)

SwiGLU(x,W,V,b,c)=SiLU(xW+b)⊗(xV+c)=x1⊗x2⊗σ(x1)\begin{align*}\text{SwiGLU}\left(x, W, V, b, c\right)&=\text{SiLU}\left(xW+b\right)\otimes\left(xV+c\right)\\&=x1\otimes x2\otimes \sigma(x1)\end{align*}

W,V∈Ric×ocW, V\in R^{ic\times oc}是权重,bb是 bias,x1=xW+b, x2=xV+cx1=xW+b, ~x2=xV+c。SiLU\text{SiLU}是 SigmoidWeighted Linear Unit, σ\sigma是 sigmoid 激活函数。

加入一个平滑因子:

SwiGLU(x,W,V,b,c)=(x1⊙s)⊗(x2⊙s)⊗σ′(x1⊙s)\text{SwiGLU}(x, W, V, b, c) = (x1 \odot s) \otimes (x2 \odot s) \otimes \sigma'(x1 \odot s)

然后将因子合并到权重中并且进行量化:

SwiGLU(x)=𝒬(𝒬(x)𝒬(W′)+b′)⊗𝒬(𝒬(x)𝒬(V′)+c′)⊗σ′(𝒬(𝒬(x)𝒬(W′)+b′))\text{SwiGLU}(x) = \mathcal{Q}(\mathcal{Q}(x)\mathcal{Q}(W') + b') \otimes \mathcal{Q}(\mathcal{Q}(x)\mathcal{Q}(V') + c') \otimes\sigma'(\mathcal{Q}(\mathcal{Q}(x)\mathcal{Q}(W') + b'))

3.3. Dynamic Integer-only MatMul (DI-MatMul)

之前的用移位 + 低精度乘法实现的总体来讲还是有很多的误差,文章提出一种新的矩阵乘法 pipeline,写作:

𝐘𝐈,sy,zpI=ℳ(s1,zp1I,s2,zp2I,𝐗𝟐𝐈)\bold{Y^I}, s_y, zp^I=\mathcal{M}(s_1, zp_1^I, s_2, zp_2^I, \bold{X_2^I})

s1,s2,sys_1, s_2, s_y是浮点数标量,表示输入、输出和量化步长;zp1I,zp2Izp_1^I, zp_2^I是整数的零点,然后将前面的浮点数改写成mI/2kIm^I/2^{k^I}的形式,就有:

YI,myI,kyI,zpI=ℳ(m1I,k1I,zp1I,X1I,m2I,k2I,zp2I,X2I)Y^I, m_y^I, k_y^I, zp^I = \mathcal{M}(m_1^I, k_1^I, zp_1^I, X_1^I, m_2^I, k_2^I, zp_2^I, X_2^I)

一次矩阵乘法操作中,

PI=(X1I−zp1I)(X2I−zp2I),Y=PIm1Im2I2k1I+k2IP^I = (X_1^I - zp_1^I)(X_2^I - zp_2^I), \quad Y = P^I \frac{m_1^I m_2^I}{2^{k_1^I + k_2^I}}

近似:

myI2kyI≈sy=(pmax⁡I−pmin⁡I)⋅m1I⋅m2I(2n−1)⋅2k1I+k2I\frac{m_y^I}{2^{k_y^I}} \approx s_y = \frac{(p_{\max}^I - p_{\min}^I) \cdot m_1^I \cdot m_2^I}{(2^n - 1) \cdot 2^{k_1^I + k_2^I}}

为了求得输出的量化尺度,有:

argminmyI,kyI⁡‖(pmax⁡I−pmin⁡I)⋅m1I⋅m2I(2n−1)⋅2k1I+k2I−myI2kyI‖1s.t.myI,kyI∈{0,1,2,…,255}\text{arg} \min_{m_y^I, k_y^I} \left\| \frac{(p_{\max}^I - p_{\min}^I) \cdot m_1^I \cdot m_2^I}{(2^n - 1) \cdot 2^{k_1^I + k_2^I}} - \frac{m_y^I}{2^{k_y^I}} \right\|_1 \quad \text{s.t.} \quad m_y^I, k_y^I \in \{0, 1, 2, \ldots, 255\}

求解其他的有:

kyI=⌊log2⁡(256⋅(2n−1)⋅2k1I+k2Ipmax⁡I−pmin⁡I)⌋myI=((pmax⁡I−pmin⁡I)⋅m1I⋅m2In1)≫(k1I+k2I−kyI)zyI=−pmin⁡I⋅(2n−1)pmax⁡I−pmin⁡I,YI=⌊PI−pmin⁡Ipmax⁡I−pmin⁡I⋅(2n−1)⌋k_y^I = \left\lfloor \log_2 \left( \frac{256 \cdot (2^n - 1) \cdot 2^{k_1^I + k_2^I}}{p_{\max}^I - p_{\min}^I} \right) \right\rfloor \\ m_y^I = \left( \frac{(p_{\max}^I - p_{\min}^I) \cdot m_1^I \cdot m_2^I}{n_1} \right) \gg (k_1^I + k_2^I - k_y^I) \\ z_y^I = - \frac{p_{\min}^I \cdot (2^n - 1)}{p_{\max}^I - p_{\min}^I} ,\quad Y^I = \left\lfloor \frac{P^I - p_{\min}^I}{p_{\max}^I - p_{\min}^I} \cdot (2^n - 1) \right\rfloor

感觉就是更复杂的 quantization,替换掉了原本的乘法 + 移位的操作

3.4. Dynamic Non-Linear Integer-only Operations

3.4.1. Dynamic Integer-only Clipped Softmax & Dynamic Integer-only Exponent function

DI-Exp:

exI⋅(mxI≫kxI)≈2xI⋅(mxI≫kxI)⋅log2⁡e=2xI⋅mxI⋅log2⁡e≫kxIe^{x^I \cdot (m_x^I \gg k_x^I)} \approx 2^{x^I \cdot (m_x^I \gg k_x^I) \cdot \log_2 e} = 2^{x^I \cdot m_x^I \cdot \log_2 e \gg k_x^I}

化简得到:

exI⋅(mxI≫kxI)≈2(1−rI2⋅sI)≫qIe^{x^I \cdot (m_x^I \gg k_x^I)} \approx 2 \left( 1 - \frac{r^I}{2 \cdot s^I} \right) \gg q^I

计算上面的pp的时候:

pmin⁡′=max⁡(pmin⁡I,pmax⁡I−cI),where cI=mcI2kcp'_{\min} = \max \left( p_{\min}^I, p_{\max}^I - c^I \right), \quad \text{where } c^I = \frac{m_c^I}{2^{k_c}}

3.4.2. Dynamic Integer-only Normalization & Dynamic Integer-only SwiGLU

DI-Norm:

Untitled.png

per channel 量化

DI-SwiGLU:

Untitled.png

4. Experiments

I-LLM 中所有的非线性的算子都是 W8A8 的,线性算子看配置,可能是 W4A8huoqitade .

4.1. Quantitative Results

Untitled.png

Untitled.png

Untitled.png

4.2. Ablation Study

Untitled.png

5. Conclusion

感觉和 I-ViT、I-BERT 几个工作挺像的,提出来的各种 DI-OP 感觉没有看到特别新的东西。

文章里很多 Integer 都打成 Interger 了。