http://arxiv.org/abs/2310.16795
MLSys 2024

1.背景和动机

为了解决大型模型的高推理成本问题，MoE架构被提出。MoE通过稀疏路由的方式，将输入分配给多个专家（experts）中的一小部分，以实现更快的推理速度和更高的模型质量。但这种架构也带来了巨大的参数量，例如SwitchTransformer-c2048模型就有1.6万亿参数。MoE模型的参数量巨大，需要数TB级的存储空间，这使得它们在实际部署时面临内存和成本的挑战，尤其是在需要大规模并行计算的场合。
为了降低MoE模型的内存和存储需求，同时保持模型性能，模型压缩成为了一个重要的研究方向。传统的压缩技术，如量化和稀疏性，虽然在一定程度上有效，但对于参数量达到万亿级别的模型来说，仍然不足以实现高效的压缩。
本文提出了QMoE，一种新的压缩和执行框架，旨在实现对万亿参数MoE模型的高效压缩和推理。QMoE通过设计一种可扩展的算法，将模型压缩到每个参数不到1比特的大小，并与定制的GPU解码内核协同设计，以实现端到端的高效压缩推理，且运行时开销相对较小。

http://arxiv.org/abs/2310.16795
Fig 1. 量化结果

http://arxiv.org/abs/2310.16795


作者首先考虑了Huffman和LZW两种常用于文件压缩的方法。但是Huffman方法的解码依赖于上文已经被解析的参数，并行性低；且变长的编码方式在实现上和存储的时候也是较为困难的。作者总结出了MoE量化的4个难点：

  1. 现有的压缩方法，如量化和稀疏性，通常只能在不显著损失精度的情况下将模型的精度降低到每个参数3或4比特，或者达到大约50%的稀疏度。然而，要使万亿参数的MoE模型实用化，需要比16位精度高出10到20倍的压缩率，即平均每个参数少于1比特。
  2. 将现有的压缩方法应用于比大型dense模型大一个数量级的MoE模型时，会遇到内存、性能和可靠性方面的障碍。MoE模型由于其稀疏性，需要处理的内存和数据量巨大。即量化过程需要的内存太大，且可能会出现因为corner case导致量化失败的问题。
  3. 实现每个参数少于1比特的压缩率需要一个非平凡的自定义压缩格式，并且这种格式需要配备在GPU等加速器上高效的解码算法，以避免在压缩模型上进行推理时出现重大的处理延迟（比如要避免Huffman方法的同步）。
  4. 为了应对上述挑战，需要在系统级别进行设计和优化，包括优化激活卸载、使用列表缓冲区来支持样本访问、延迟权重获取以减少内存占用、专家分组以提高GPU利用率，以及进行鲁棒性修改以处理在压缩具有数万个层的模型时可能遇到的罕见corner case。

2. 算法

2.1 使用GPTQ量化

http://arxiv.org/abs/2310.16795
Fig 2. 使用GPTQ量化流程

http://arxiv.org/abs/2310.16795


具体来说，我们维护一个大型缓冲区$B$，并按以下方式更新 Transformer 块的Dense部分：

  1. 从CPU到GPU抓取一个 “样本” $X$，其中包含数百个Token
  2. 通过相应的Dense Layer，得到结果$Y$
  3. 计算并存储$Y$中标记的专家分配
  4. 将$Y$送回CPU并覆盖$B$中的$X$

并且对于稀疏部分，分别对专家进行循环：

  1. 从CPU到GPU获取$B$中所有被分配给专家$E$的单独Token，记作$X_{E}$
  2. 使用它们来生成压缩后的专家$E^{’}$（例如，使用GPTQ算法）
  3. 通过$E^{’}$模块以获得$Y_{E^{’}}$
  4. 将$Y_{E^{’}}$发送回CPU，并在B中覆盖$X_{E}$

作者在这里还引入了List Buffering、Lazy Weight Fetching和Expert Grouping技巧

2.1.1 List Buffering

为了有效地支持对Dense模型的访问，以及对专家tokens的完全向量化查询，我们将$B$存储为列表缓冲数据结构。这可以被看作是一个包含所有tokens隐藏状态的巨大连续缓冲区，以及分隔符索引，这些索引标志着各个样本之间的边界。下图展示了这种存储格式。这种数据结构对效率至关重要；对于大量样本计数，通过掩码迭代样本并获取相关tokens的方法是很慢的，而作者提出的方法则有大幅度改进。

http://arxiv.org/abs/2310.16795
Fig 3. list buffering

http://arxiv.org/abs/2310.16795

2.1.2 Lazy Weight Fetching

由于1.6万亿参数模型的权重占用了超过3TB的存储空间，它们甚至无法存储在CPU的RAM中。因此，我们按需直接从磁盘存储中懒加载它们。按照推理的流程，我们需要将所有的参数从磁盘搬移到内存中完整的一整次。

2.1.3 Experts Grouping

此外，为了避免GPU的利用率不足，作者将多个专家组合在一起，并应用GPTQ算法的联合批处理变体。

2.2 字典生成

对于量化后得到的Ternary Pair ${w_{min}, 0, w_{max}}$，在很多的情况下，是0居多的，也就是说是稀疏的，那么对于稀疏矩阵可以用CSR等方法来存储。但是使用传统的稀疏矩阵存储方法压缩比还是不够，作者团队使用了一种更加偏向于文件压缩的思路来进行量化后的参数压缩，这个方法就使用到了字典查找的方法。字典查找的方法还是比较通俗易懂的，以下面的例子来举例：

对于“001002003…”我们可以统计该串里面的子串的出现频率，比如001，002，003出现的频率高，那么我们可以将他们编码成 A,B,C然后仅需要三个char的空间“ABC”就可以表示一个压缩后的文件。

使用字典来压缩就是本文为什么可以做到平均每个weight小于1bit量化的来源。

一般来说，假设三元权重矩阵（用 0、1、2 表示）的值接近独立分布，其分布如下:

$$ P(0)=p_{0}, P(1)=P(2)=\frac{1-p_{0}}{2} $$

http://arxiv.org/abs/2310.16795
Fig 4. 字典生成算法

http://arxiv.org/abs/2310.16795

然后，作者用上图中的算法来生成字典。字典的Key是16bit的类型，字典的Value是一个长度在14以下的非空对数组(即：[(t1, t2), (t0, t1), …])。本文在生成字典的时候做了两个限制：1. 字典的Key是16位的 2. 字典的Value最多是14个Pair。

该算法的每次循环，会找出队列中出现概率最高的Value，然后对其Value Concat上一个新的Pair，这样做的好处是：可以让Cache Locality特性发挥的很好，我们会在2.3小节中看到。 作者一共生成了有$2^{16}$个条目的字典。

字典中的Value是像下图所示排布的：

http://arxiv.org/abs/2310.16795
Fig 5. Value的排布

http://arxiv.org/abs/2310.16795

一个字典中的Value包含2个int32类型，之所以不用int64是出于GPU访存冲突的考量。每个int32的4个低位表示该Value有多少个有效的pairs。每2个2bits表示一个pair，每个2bits表示一个weight。

2.3 Decode

在Decode阶段，我们需要知道当前wrap和thread的ID，然后用shift操作符提取到对应字典中的元素。核心伪代码如下：

int enc = w_comp_block[warp][j];
int wx14 = dec[2 * enc + (lane / 14)]; 
int ter = (wx14 >> (4 + 2 * (lane % 14))) & 0x3; 
float w = deq[ter][thread]; 31 res += w * x_shared[idx + lane]; 
idx += 2 * (wx14 & 0xf);

其中w_comp_block是已经压缩过的weight；dec是一维字典；deq是预先提取出来的三元组；x_shared是输入X。作者在这里融合了解码和点乘的操作。

3. 总结与思考

QMoE有着非常好的效果，但是现在的MoE模型普遍不是开源的，并且有特定的Infra设施。QMoE的进一步验证和实施还需要更多的业内工程跟进。
不理解的地方：生成字典的Algorithm1为什么可以生成所有情况的Value呢？会有不同的组合情况产生吗？