这篇文章提出了如何在资源受限的消费级硬件上高效地运行稀疏专家混合（MoE）语言模型的方法。将Mixtral-8x7B这个需要100G以上算力才能部署的模型在12G显存+11G内存的组合下跑出来。
论文：https://arxiv.org/abs/2312.17238
Colab代码：https://colab.research.google.com/drive/1ZkC0k487oBEF19R8_9nq2MSHFyQ6OspG?usp=drive_link

引言与背景

论文的引言部分介绍了大规模预训练语言模型（LLMs）在自然语言处理领域的重要性。这些模型如GPT-3、GPT-4以及其他开放访问的LLMs（如LLaMA、Falcon、BLOOM等）推动了语言技术的迅猛发展。然而，LLMs 的庞大参数量使得它们的推理成本极高，通常需要高端的GPU设备才能运行，限制了它们在普通硬件上的使用。
为了缓解这个问题，稀疏的专家混合（MoE）模型被提出。MoE通过只激活模型中的一部分“专家”来计算每个输入，从而提高了计算效率。然而，MoE模型的规模依然庞大，尤其是在需要多GPU的环境下。因此，如何在消费级硬件上运行这些模型是一个重要的研究问题。

三大解决策略

Mixtral-8x7B模型中的总参数为46.7亿，专家构成45.1亿（96.6%），在内存受限的情况下，减少专家切换时GPU与RAM之间的数据传输对MoE模型进行推理很关键。作者主要提出通过LRU缓存（LRU caching）和专家的推测性提前加载（Speculative Expert Loading）来减少GPU与RAM之间的数据传输，从而加速推理过程。关键创新点包括：

  1. LRU缓存专家重用模式：MoE模型在处理连续的token时，某些专家会被频繁地重用。因此，作者设计了一种LRU缓存机制，利用这种专家重用的规律来减少GPU-RAM之间的通信开销。

  2. 推测性专家加载：由于推理过程中无法提前确定下一层需要加载的专家，因此作者提出了一种基于推测的加载机制，通过对前一层的隐藏状态应用下一层的门控函数来猜测即将需要的专家（可能是因为有residual的原因）。这种机制在推测正确时，下一层的计算可以立即开始，显著减少了推理延迟。

  3. 混合量化技术：在模型压缩方面，作者使用了一种半二次量化（Half Quadratic Quantization, HQQ）的方法对专家层进行更高的压缩，同时保持其他层的较高精度。这种量化策略有效减少了模型大小，并保持了较好的推理性能。

实验与评估

模型

用到的模型：Mixtral 8x7B ，一个主流的MOE模型，在大多数基准测试中优于或等价于Llama2 70B, GPT3.5，且推理速度比Llama2 70B快六倍！
Mixtral 8x7B 是decoder-only model, 其中 FFN 从8个不同的参数组（专家）中进行挑选，在每一层，每个token, router network 都会选2个组来进行生成并进行组合：

  1. 支持32K上下文
  2. 支持英语，法语，意大利语，德语，西班牙语（中文支持很差）
  3. 在代码生成上很强
  4. 能被微调成一个高分（MT-Bench）的 instruction-following model

评估结论

论文在不同硬件配置（如RTX 3060、T4等）下对提出的方法进行了详尽的实验评估，得出了以下几个主要结论：

  1. 专家缓存与推测加载的有效性：通过测试不同缓存大小和提前加载的专家数量，作者发现缓存命中率和推测加载的准确率显著提高了模型的推理速度。例如，在缓存大小为4时，缓存命中率可以达到约0.8，推测加载大小为2时，推测加载的准确率则可达到0.9以上。

  2. 量化对模型性能的影响：通过对模型进行不同量化比特的测试，作者验证了在保持较好准确率的同时，量化可以有效减少模型大小。例如，使用2-bit量化时，模型的推理延迟显著降低，同时在WikiText2和C4数据集上的困惑度仅略有上升。

  3. 实际推理性能：在使用完整的算法时，消费级硬件上如RTX 3060和T4可以达到每秒生成2-3个token的性能，远远优于直接在设备内存不足的情况下推理时的性能表现。

结论与未来工作

论文总结了该方法在推理速度上相较于传统的加载方式有显著提高，尤其是在消费级硬件和免费云平台（如Google Colab）上，使得大规模稀疏MoE模型的使用更加广泛化。未来的研究方向可能包括进一步优化专家预测加载算法，探索其他的推理加速方法。

总结

这篇论文解决了大规模稀疏专家模型在推理时的硬件瓶颈问题，提出了一种通过专家缓存与预测加载来优化推理速度的方案，并使用混合量化技术在保证准确率的同时大幅减少了模型大小和推理时间。对于希望在低端硬件上使用大规模语言模型的研究人员来说，本文的贡献提供了一个具有实用价值的解决方案。

Ideas

  1. “Note that out of 46.7B total parameters in the Mixtral-8x7B model, the experts constitute 45.1B (96.6%).” 专家参数占了主导位置，这种异构型能否用于边缘-云端计算，隐私保护等
  2. 对于多用户多对话的在线推理服务系统，采取批量、并行的策略合理使用experts参数，增加吞吐、降低延迟等
  3. 将这种offload方法引入到训练过程中，可以显著扩大模型或数据集的规模