LLM2D
Quamba2:选择状态空间模型的鲁棒且可扩展的后训练量化框架
Quamba2: A Robust and Scalable Post-training Quantization Framework for Selective State Space Models
作者: Hung-Yueh Chiang, Chi-Chih Chang, Natalia Frumkin, Kai-Chiang Wu, Mohamed S. Abdelfattah, Diana Marculescu
发布日期: 4/4/2025
arXiv ID: oai:arXiv.org:2503.22879v2

摘要

arXiv:2503.22879v2 宣告类型: replace-cross 摘要:状态空间模型(SSMs)因其一致的内存使用和高性能,正在成为替代变压器的有吸引力的选择。尽管如此,由于存储需求和计算能力的限制,将SSMs扩展到云服务或有限资源的设备上依然是具有挑战性的。为了克服这个问题,使用低位宽数据格式对SSMs进行量化可以减小模型大小并从硬件加速中受益。然而,由于SSMs容易产生量化引起的误差,最近的努力集中在在不牺牲性能的情况下优化特定的模型或位宽以提高效率。然而,不同的位宽配置对于不同的场景是必要的,例如W4A8用于增强大批次解码速度,而W4A16则用于增强单用户短提示生成速度。为了应对这一挑战,我们提出了Quamba2,它兼容Mamba1和Mamba2的W8A8、W4A8和W4A16配置,以应对SSM在各种平台上的部署需求。基于状态空间模型的通道顺序保持和激活持久性,我们提出了一种分批离线量化方法,将线性递归的输入按输入$x$排序和聚类,结合对输入依赖参数$B$和$C$的按状态组量化。为确保状态空间模型输出的计算不变性,我们根据聚类序列离线重新安排权重。实验结果显示,Quamba2-8B优于几种最新的SSM量化方法,在预填充和生成阶段分别提供1.3倍和3倍的速度提升,同时实现4倍的内存减少,平均准确率下降仅为1.6%。我们在MMLU上的评估显示了我们框架的通用性和鲁棒性。代码和量化模型将发布在:https://github.com/enyac-group/Quamba。