Quamba2实现Mamba模型多比特量化,兼顾速度与内存压缩。
Quamba2: A Robust and Scalable Post-training Quantization Framework for Selective State Space Models
- 通过输入排序聚类和状态组量化,实现8位以下高效压缩。
- 在预填充和生成阶段分别提速1.3倍和3倍,内存减少4倍。
- 支持W8A8/W4A8/W4A16多种配置,适合不同部署场景。
状态空间模型(SSMs)因其稳定的内存占用和高性能,正成为Transformer的有力替代。然而,在云服务或资源受限设备上扩展SSMs仍面临存储和计算挑战。通过低比特量化可减小模型体积并利用硬件加速。由于SSMs易受量化误差影响,现有工作多针对特定模型或比特位宽优化性能。但不同场景需不同配置:如W4A8提升大批次解码速度,W4A16则利于单用户短提示生成。为此,我们提出Quamba2,兼容Mamba1和Mamba2骨干网络的W8A8、W4A8和W4A16配置,满足多样部署需求。基于SSMs的通道顺序保持与激活持久性,我们提出一种离线方法:对线性递归输入x进行8位量化,通过排序与聚类实现;对依赖输入的参数B、C采用逐状态组量化。为保证输出计算不变性,离线重排权重以匹配聚类序列。实验表明,Quamba2-8B优于两项先进量化方法,在预填充和生成阶段分别实现1.3倍和3倍加速,内存减少4倍,仅带来1.6%平均精度下降。MMLU评估验证了框架的泛化性和鲁棒性。代码与量化模型将开源于https://github.com/enyac-group/Quamba。
原文摘要 · Abstract (English)
State Space Models (SSMs) are emerging as a compelling alternative to Transformers because of their consistent memory usage and high performance. Despite this, scaling up SSMs on cloud services or limited-resource devices is challenging due to their storage requirements and computational power. To overcome this, quantizing SSMs with low bit-width data formats can reduce model size and benefit from hardware acceleration. As SSMs are prone to quantization-induced errors, recent efforts have focused on optimizing a particular model or bit-width for efficiency without sacrificing performance. However, distinct bit-width configurations are essential for different scenarios, like W4A8 for boosting large-batch decoding speed, and W4A16 for enhancing generation speed in short prompt applications for a single user. To this end, we present Quamba2, compatible with W8A8, W4A8, and W4A16 for both Mamba1 and Mamba2 backbones, addressing the growing demand for SSM deployment on various platforms. Based on the channel order preserving and activation persistence of SSMs, we propose an offline approach to quantize inputs of a linear recurrence in 8-bit by sorting and clustering for input $x$, combined with a per-state-group quantization for input-dependent parameters $B$ and $C$. To ensure compute-invariance in the SSM output, we rearrange weights offline according to the clustering sequence. The experiments show that Quamba2-8B outperforms two state-of-the-art SSM quantization methods and delivers 1.3$\times$ and 3$\times$ speed-ups in the pre-filling and generation stages, respectively, while offering 4$\times$ memory reduction with only a $1.6\%$ average accuracy drop. The evaluation on MMLU shows the generalizability and robustness of our framework. The code and quantized models will be released at: https://github.com/enyac-group/Quamba.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。