AI论文精选
⏶7

BlockFFN: 面向端侧加速友好的块粒度激活稀疏专家混合模型

发表
由 Chenyang SongChenyang Song 提交
作者: Chenyang SongChenyang Song, Weilin Zhao, Xu Han, Chaojun Xiao, Yingfa Chen, Yuxuan Li, Zhiyuan Liu, Maosong Sun

摘要

为缓解大语言模型（LLM）的计算负担，以专家混合（MoE）为代表的激活稀疏架构受到了越来越多的关注。然而，传统MoE的不可微和不灵活的路由损害了模型性能。此外，虽然每个词元仅激活少量参数，但这些稀疏激活的架构表现出较低的块级稀疏性，这表明多个连续词元的联合激活了大量比例的参数。这种稀疏模式在低资源条件（例如，端侧设备）下不利于加速，并且与主流加速技术（例如，推测解码）不兼容。为了应对这些挑战，我们引入了一种新型MoE架构BlockFFN，以及其高效的训练和部署技术。具体来说，我们使用了一个集成ReLU激活和RMSNorm的路由器，以实现可微和灵活的路由。接下来，为了同时提升词元级稀疏性（TLS）和块级稀疏性（CLS），我们设计了CLS感知训练目标，使BlockFFN更利于加速。最后，我们实现了高效的加速核，首次结合了激活稀疏性和推测解码。实验结果表明，BlockFFN在性能上优于其他MoE基线，实现了超过80%的词元级稀疏性（TLS）和70%的8词元块级稀疏性（CLS）。我们的核在真实端侧设备上比密集模型实现了高达3.67倍的加速。所有代码和检查点均已公开（https://github.com/thunlp/BlockFFN）。
查看 arXiv 页面查看 PDF

评论

Chenyang SongChenyang Song
论文作者
论文提交者

在本文中，我们主要解决现有MoE架构面临的两个挑战：

  1. 性能折衷，由不完善的路由引起，尤其是传统路由范式的不可微性和不灵活性问题；

  2. 加速不友好，由低块级稀疏性（CLS）引起，尤其是在同时处理多个token的条件下，例如卸载和推测解码。

为了应对上述挑战，我们引入了BlockFFN，一种新颖的MoE架构，以及其训练技术和高效的端侧部署。

  1. 在模型架构方面，我们提出了BlockFFN，一种新颖的MoE范式，通过路由器模块将性能折衷降至最低，并结合了ReLU激活和RMSNorm。通过实验，我们证明了它比其他MoE基线（如TopK、DeepSeekMoE、GRIN和ReMoE）具有更好的性能。

  2. 在训练技术方面，我们引入了CLS感知训练目标，以提高BlockFFN的CLS以及传统的token级稀疏性（TLS）。在实验中，我们获得了高于80%的平均TLS值和高于70%的8-token CLS值。

  3. 在端侧部署方面，我们为BlockFFN实现了高效的加速内核，首次结合了激活稀疏性和推测解码。在NVIDIA Jetson Orin NX上，与基线自回归（AR）解码相比，该内核实现了3.67倍的加速比。