8张3090，1天压缩万亿参数大模型！3.2TB骤降至160GB，压缩率高达20倍

文章来源：新智元

最近，来自ISTA的研究人员提出了一种全新的模型量化方法QMoE，可以将1.6万亿个参数的SwitchTransformer压缩到160GB以下（每个参数0.8位），且精度损失很小。

图片来源：由无界 AI生成

随着GPT-4的架构被知名业内大佬「开源」，混合专家架构（MoE）再次成为了研究的重点。

GPT-4拥有16个专家模型，总共包含1.8万亿个参数。每生成一个token需要使用大约2800亿参数和560TFLOPs

然而，模型更快、更准确的代价，则是巨大的参数量，和随之而来的高昂成本。

比如，1.6万亿参数的SwitchTransformer-c2048模型，需要3.2TB的GPU显存才能有效运行。

为了解决这一问题，来自奥地利科技学院（ISTA）的研究人员提出了一种全新的压缩和执行框架——QMoE。

论文地址：https://arxiv.org/abs/2310.16795

通过采用专门设计的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块的稠密部分进行更新：

1. 从CPU到GPU，抓取一个包含有几百个token的「样本」X；

2. 通过对应的稠密层，得到结果Y；

3. 计算并存储Y中token的专家分配；

4. 将Y发送回CPU并覆盖B中的X。

对于稀疏部分：

1. 从CPU到GPU，抓取B中所有已分配给专家E的token，用X_E表示。

2. 利用它们生成压缩的专家E'（例如，使用GPTQ）。

3. 通过E'运行X_E，得到Y_E'。

4. 将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的高效性。