文章

E³-PRUNER

llm-inference
本文目录 7

E³-PRUNER 论文阅读

华为的一个剪枝方法

简单来说，逻辑上有两个关键点：

  1. 剪枝
  2. 微调模型

在训练的循环过程中，在开始的循环中先联合了剪枝和微调，之后只进行微调

搜索剪枝

Naive的来看，给定一个参数KK，让模型最终剪去KK层的layer，只需要暴力遍历所有的结果并评估分数就行。只要对所有的数据集进行C(L,K)=L!K!(L−K)!C(L,K)=\frac{L!}{K!(L-K)!}次计算

抛开这千万到万亿次遍历×\times万级的数据的计算，直接剪枝还容易导致出现伪最优，因为一次性剪去KK层之后，输出的肯定是胡话，在胡话里选出一个最优的是不合适的

那如果一个一个的剪呢，也同样容易造成错误。因为剪去一层是不可复原的，也就是剪的这一个可能实际上是一个比较有用的层，但是在当前的数据中没有显现；也就是说贪心算法在这个过程中基本上得到的不是最优的结果

综上都是离散剪枝带来的问题。离散剪枝导致没有容错，而E3-PRUNER\text{E}^3\text{-PRUNER} 使用了 Gumbel-TopKGumbel\text{-}TopK 方法，旨在将离散的 Top-KTop\text{-}K 变成可微分的

Gumbel-TopK 算法

Require: Layer scoreSS, temperatureτ\tau, number of retained layers kk.

1:S1=S+Gumbel(0,1)2:M~1=03:for  i=1  to  k  do4:Si+1=Si+log⁡(1−Softmax(Siτ))5:M~i+1=M~i+Softmax(Siτ)6:end for7:M=STE(TopK(M~k+1,k))8:return  Pruning mask  M\begin{aligned} &1: \quad S_1 = S + \text{Gumbel}(0,1) \\ &2: \quad \tilde{M}_1 = 0 \\ &3: \quad \textbf{for} \; i = 1 \; \textbf{to} \; k \; \textbf{do} \\ &4: \quad\quad S_{i+1} = S_i + \log(1 - \text{Softmax}(\frac{S_i}{\tau})) \\ &5: \quad\quad \tilde{M}_{i+1} = \tilde{M}_i + \text{Softmax}(\frac{S_i}{\tau}) \\ &6: \quad \textbf{end for} \\ &7: \quad M = \text{STE}(\text{TopK}(\tilde{M}_{k+1}, k)) \\ &8: \quad \textbf{return} \; \text{Pruning mask} \; M \end{aligned}

如上面的公式展示的，每一层都有一个分数SS，在初始的时候给每个分数都加了一个GumbelGumbel噪声，防止分数平均导致微小变化产生巨大分数变化

第4行会对重复的高分进行”惩罚”，防止重复选几个特定的层；关键在于STE(Straight-Through Estimator，直通估计器)让back-propagation能够运算

至于τ\tau，则采用的是线性退火的形式，在初期能探索更多的可能

迭代过程如下：

M~i+1=M~i+Softmax(Siτ),Si+1=Si+log⁡(1−Softmax(Siτ))\tilde{M}_{i+1} = \tilde{M}_i + \text{Softmax}(\frac{S_i}{\tau}), \quad S_{i+1} = S_i + \log(1 - \text{Softmax}(\frac{S_i}{\tau}))

微调

1.联合训练

剪去某个特殊层之后，模型的权重应该跟着微调，否则就会精度断崖下降

算法1：

Params:

  • LLM Θ\Theta
  • total iteration TT
  • mask searching iteration turns TMT_M
  • layer score SS
  • layer mask MM
  • total layers LL.
1:for  t=1  to  T  do2:τ=1−β⋅t/T,  k′=L−⌈(L−k)⋅t/T⌉3:if  t≤TM  :4:M=Gumbel-TopK(S,τ,k′)5:S=S−λ⋅∂L(x,Θ,M)/∂S6:end if7:Θ=Θ−λ⋅∂L(x,Θ,M)/∂Θ8:end for9:return  Pruned model  Θ∗\begin{aligned} &1: \quad \textbf{for} \; t = 1 \; \textbf{to} \; T \; \textbf{do} \\ &2: \quad\quad \tau = 1 - \beta \cdot t/T, \; k' = L - \lceil (L - k) \cdot t/T \rceil \\ &3: \quad\quad \textbf{if} \; t \leq T_M \; \textbf{:} \\ &4: \quad\quad\quad M = \text{Gumbel-TopK}(S, \tau, k') \\ &5: \quad\quad\quad S = S - \lambda \cdot \partial \mathcal{L}(x, \Theta, M) / \partial S \\ &6: \quad\quad \textbf{end if} \\ &7: \quad\quad \Theta = \Theta - \lambda \cdot \partial \mathcal{L}(x, \Theta, M) / \partial \Theta \\ &8: \quad \textbf{end for} \\ &9: \quad \textbf{return} \; \text{Pruned model} \; \Theta^* \end{aligned}

这里值得关注的就是这里的反向传播，得益于 STE (Straight-Through Estimator) 得以实现

而这里的 Gumbel-TopK 的算法如下

算法2:

Params:

  • Layer score SS
  • temperature τ\tau
  • layers kk
1:S1=S+Gumbel(0,1)2:M1~=03:for i=1 to k do4:M=Gumbel-TopK(S,τ,k′)5:S=S−λ⋅∂L(x,Θ,M)/∂S6:end if7:Θ=Θ−λ⋅∂L(x,Θ,M)/∂Θ8:end for9:return  Pruned model  Θ∗\begin{aligned} &1: \quad S_1 = S + \text{Gumbel}(0,1) \\ &2: \quad \tilde{M_1} = 0 \\ &3: \quad \text{for } i = 1 \text{ to } k \text{ do}\\ &4: \quad\quad\quad M = \text{Gumbel-TopK}(S, \tau, k') \\ &5: \quad\quad\quad S = S - \lambda \cdot \partial \mathcal{L}(x, \Theta, M) / \partial S \\ &6: \quad\quad \textbf{end if} \\ &7: \quad\quad \Theta = \Theta - \lambda \cdot \partial \mathcal{L}(x, \Theta, M) / \partial \Theta \\ &8: \quad \textbf{end for} \\ &9: \quad \textbf{return} \; \text{Pruned model} \; \Theta^* \end{aligned}

文中提到的loss函数是这样的：

{Θ∗,M∗}=arg⁡min⁡Θ,ML(x;Θ,M)s.t.  ∥M∥0=k\left\{\Theta^{*}, M^{*}\right\} = \arg\min_{\Theta, M} \mathcal{L}(x ; \Theta, M) \quad \text{s.t.} \; \|M\|_0 = k

也就是自适应蒸馏损失：

Ladaptive-KD=1N∑i=1NH(Softmax(zt(i))c)⋅∑c∈SiKSoftmax(zt(i))c⋅log⁡(Softmax(zt(i))cSoftmax(zs(i))c)\mathcal{L}_{\text{adaptive-KD}} = \frac{1}{N} \sum_{i=1}^{N} \mathcal{H}\left(\text{Softmax}\left(z_t^{(i)}\right)_c\right) \cdot \sum_{c \in S_i^K} \text{Softmax}\left(z_t^{(i)}\right)_c \cdot \log\left(\frac{\text{Softmax}\left(z_t^{(i)}\right)_c}{\text{Softmax}\left(z_s^{(i)}\right)_c}\right)

它由两部分组成：

token损失：

Ltoken=∑c∈SiKSoftmax(zt(i))c⋅log⁡(Softmax(zt(i))cSoftmax(zs(i))c)\mathcal{L}_{\text{token}} = \sum_{c \in S_i^K} \text{Softmax}\left(z_t^{(i)}\right)_c \cdot \log\left(\frac{\text{Softmax}\left(z_t^{(i)}\right)_c}{\text{Softmax}\left(z_s^{(i)}\right)_c}\right)

这里的zsz_s是 student 的输出，ztz_t是 teacher 的输出。由于现场推理开销大，所以 teacher 输出的 top-K 个 logits 是事先计算好并保存的

权重：

Wi=H(Softmax(zt(i))c)\mathcal{W}_i = \mathcal{H}\left(\text{Softmax}\left(z_t^{(i)}\right)_c\right)

2.后训练

如算法1展示的那样，在kk轮的联合训练之后，再进行了T−kT-k轮的单独对模型的调参，目标是在固定模型框架的基础上恢复精度

延申

比起直接对齐向量方向的训练，对齐输出分布的训练往往上限更高，因为方向上的对齐不一定能带来最终logits的更加准确

当大模型最终输出了一个h之后，由z=Wuhz = W_{u}h得到最终的logits

由于剪枝之后预测的hh多少有偏差。我们假设是h′=h+Δhh'=h+\Delta h，这时z′=Wuh′=Wu(h+Δh)=z+WuΔhz' = W_uh' = W_u (h + \Delta h) = z + W_u \Delta h

可以看到最终logits的分布并不是单纯依赖于Δh\Delta h也就是向量的相似，而是依赖于WuΔhW_u \Delta h这个整体

也就是说即使在Δh\Delta h相差很大的情况下，只要这个Δh\Delta h向量中的元素的分布能和WuW_u矩阵的零空间高度重合，最终的logits就是被很好拟合的

反之，就算得到很小的Δh\Delta h，如果他恰巧和WuW_u中的有大元素的空间重合，产生的Δz\Delta z这个误差也会很大，这就是为什么方向的对齐的上限会低一点

但是相应的他的训练成本会比输出对齐的训练要低，最后得到的结果和输出分布训练的向量差距也不会那么大，只要post train的时候去对齐一下输出，理论上就能得到更好的效果