当前位置:AIGC资讯 > AIGC > 正文

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

随着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的高效性。

从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。

更新时间 2023-10-31