【新智元导读】 最近，来自ISTA的研究人员提出了一种全新的模型量化方法QMoE，可以将1.6万亿个参数的SwitchTransformer压缩到160GB以下（每个参数0.8位），且精度损失很小。
随着GPT-4的架构被知名业内大佬「开源」，混合专家架构（MoE）再次成为了研究的重点。
GPT-4拥有16个专家模型，总共包含1.8万亿个参数。每生成一个token需要使用大约2800亿参数和560TFLOPs
然而，模型更快、更准确的代价，则是巨大的参数量，和随之而来的高昂成本。
比如，1.6万亿参数的SwitchTransformer-c2048模型，需要3.2TB的GPU显存才能有效运行。
为了解决这一问题，来自奥地利科技学院（ISTA）的研究人员提出了一种全新的压缩和执行框架------QMoE。
通过采用专门设计的GPU解码内核，QMoE具备了高效的端到端压缩推理------不仅可以实现高达20倍的压缩率，而且只会产生轻微的精度损失。
具体而言，QMoE仅需单个GPU服务器，就可以在一天内将1.6万亿参数的SwitchTransformer-c2048模型压缩至不到160GB，相当于每参数只有0.8位。
如此一来，就可以在4张英伟达RTX A6000或8张英伟达RTX 3090 GPU上运行，而推理时的开销还不到未压缩模型的5%。
MoE模型量化
混合模型（MoE）的核心理念是通过增加模型参数量，来提高网络的建模能力，同时与标准的前馈架构相比，保持计算成本几乎不变。
由于处理每个输入token时仅需调用网络中的一小部分，因此这种设计可以利用100个甚至1000个「专家」来构建超大规模的模型，并进行高效的训练和推理。
事实证明，在推理速度相当的情况下，MoE可以大幅提高准确率和训练速度。但如此庞大的体积，也就意味着需要大量的显存才能让模型跑起来。
压缩MoE的一个主要挑战是需要维持庞大的激活集。
对此，可以通过精心安排模型执行的方式，将需要计算的中间数据控制在一小部分。从而把主存储从GPU卸载到价格更便宜、数量更多的CPU内存中。
具体来说就是，维持一个大型缓冲区B，并按照以下步骤对Transformer块的稠密部分进行更新：
- 从CPU到GPU，抓取一个包含有几百个token的「样本」X；
- 通过对应的稠密层，得到结果Y；
- 计算并存储Y中token的专家分配；
- 将Y发送回CPU并覆盖B中的X。
对于稀疏部分：
-
从CPU到GPU，抓取B中所有已分配给专家E的token，用X_E表示。
-
利用它们生成压缩的专家E'（例如，使用GPTQ）。
-
通过E'运行X_E，得到Y_E'。
-
将Y_E'送回CPU并覆盖B中的X_E。
如图2所示，这个过程最小化了内存消耗和传输成本：只需一个B的副本，每个token在每个Transformer块中只被读写了两次。
更进一步的，研究人员设计了一个编码方案和一个CUDA内核，实现了每权重低于1位的压缩，并将推理的GPU执行开销降至最低。
压缩效果
精度
首先，研究人员将所有SwitchTransformer模型量化到2位和三元精度，然后评估其验证损失。
对于128个专家，默认的校准样本数为10K；对于2048个专家，默认的校准样本数为160K。同时，研究人员也测试了0.5倍和2倍的样本数。
结果显示，使用数据依赖的量化，2位模型可以在最小的损失下实现（相对于c2048，损失为1.7%），而三元精度下的损失增加也很小（相对于c2048，损失为6.7%）。
这不仅证明了所提出的先进量化方法的有效性，而且还表明极低位宽的压缩确实适用于大规模的MoE。
此外，研究人员还在来自RedPajama的arXiv、GitHub、StackExchange和Wikipedia的数据上进行了评估。
虽然校准数据中只有<0.01%来自这些网站，但压缩后的模型依然保持了几乎与核心分布相同的性能。
就校准数据而言，增加样本数量通常会略微提高性能，在三元量化时最为明显。但在此过程中也会出现一些噪声，尤其是在2位时。
压缩
测试中，研究人员同时考虑了仅MoE模块的压缩，以及相对于整个模型及其所有元数据的压缩。
仅MoE本身的而言，所有规模都实现了>16倍的压缩率，相当于每个参数的存储空间都<1位。
在c2048上，即使是包括所有未压缩的稠密层在内，整体的压缩率也达到了19.81倍，相当于每个参数0.807位，从而将检查点大小从3142GB减少到158.6GB。
此外，还可以观察到压缩率随模型大小的增加而增加，这有两个原因：
（a）自然稀疏性增加，且研究人员针对c2048优化了编码字典；
（b）层越大，权重分布越接近独立。
运行时间
最后，研究人员评估了针对不同数量的校准数据，在单个A6000 GPU上生成压缩模型所需的时间。
结果显示，较小的模型可以在一小时内压缩完成，即便是c2048也能在不到一天的时间内完成，这证实了QMoE的高效性。
从large128到c2048，运行时间的增加与大小的差异基本成正比，尽管后者使用了多16倍的样本。这是因为每个专家的样本数量保持不变，而专家规模仅略有增加。
运行结果
首先，将压缩的矩阵-向量积内核与PyTorch标准的（未压缩）bfloat16 cuBLAS内核进行直接（孤立）比较。
图 5（左）显示了压缩内核与bfloat16内核相比，在两款不同的GPU上，MoE发现矩阵形状所耗费的时间。
虽然研究人员使用的储存性能较差，但执行压缩内核所需的时间，依然比接近理想的bfloat16基线少。在特定矩阵形状下，速度最多可提高35%。
而这些操作的延迟也非常低，其中，最小的矩阵耗时<0.02毫秒，最大的耗时<0.05毫秒。
随后，研究人员在HuggingFace中，利用压缩MoE模型的实际权重，对内核进行了端到端的基准测试。
结果如图5（右）所示，压缩模型的端到端执行速度只比标准（未压缩）的慢了<5%。
尽管每层时序更快，但速度仍略有下降，这是因为编码器有时会将多个token路由到同一个专家。
目前的实现方式是，对每个token执行单独的矩阵向量乘积，而基线执行的是更高效的联合矩阵乘法。
在一些应用中，这是一个很大的瓶颈。对此，可以在内核中引入token内循环，或者在token数量较多的情况下，先进行完全解压缩，然后再执行标准的矩阵乘法。
讨论与局限性
总结而言， QMoE是一个开源的端到端压缩和推理框架，用于解决MoE在推理过程中，内存开销过大的问题。
研究人员首次证明了，像SwitchTransformer c2048这样的万亿参数模型，可以精确压缩到每个参数小于1位，压缩率接近20倍。并且，首次在单个消费级GPU服务器上，实现了此类模型的高效端到端执行。
不过，由于只有少数大规模且精确的MoE可以被公开获得，因此研究的模型集十分有限。
此外，由于其规模庞大，大多数MoE都是在不同的定制框架中训练和部署的，这就需要复杂的手动集成才能用于进一步研究。
尽管如此，研究人员还是涵盖了一些规模最大、精度最高的MoE，特别是SwitchTransformer。