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