Quamba2: A Robust and Scalable Post-training Quantization Framework for Selective State Space Models

要約

状態空間モデル（SSM）は、一貫したメモリ使用量と高いパフォーマンスにより、Transformerに代わる有力な選択肢として台頭してきている。にもかかわらず、クラウドサービスや限られたリソースのデバイス上でSSMをスケールアップすることは、そのストレージ要件と計算能力のために困難である。これを克服するために、低ビット幅のデータフォーマットでSSMを量子化することで、モデルサイズを縮小し、ハードウェアアクセラレーションの恩恵を受けることができます。SSMは量子化に起因するエラーが発生しやすいため、最近の取り組みでは、性能を犠牲にすることなく効率性を高めるために、特定のモデルやビット幅を最適化することに焦点が当てられています。しかし、W4A8は大バッチのデコード速度を向上させ、W4A16は単一ユーザー向けの短時間アプリケーションの生成速度を向上させるなど、異なるシナリオには異なるビット幅構成が不可欠である。この目的のために、我々はMamba1とMamba2の両方のバックボーン用のW8A8、W4A8、W4A16と互換性のあるQuamba2を発表し、様々なプラットフォーム上でのSSM展開の需要の高まりに対応する。SSMのチャネル次数保持と活性化持続性に基づいて、入力$x$のソートとクラスタリング、入力依存パラメータ$B$と$C$の状態グループ毎の量子化と組み合わせることにより、8ビットでの線形再帰の入力を量子化するオフラインアプローチを提案する。SSM出力の計算不変性を確保するため、クラスタリング順序に従って重みをオフラインで再配列する。実験の結果、Quamba2-8Bは、いくつかの最新のSSM量子化手法を凌駕し、プリフィリングと生成段階でそれぞれ1.3$times$と3$times$の高速化を実現し、平均精度低下$1.6%$のみで4$times$のメモリ削減を提供する。MMLUでの評価は、本フレームワークの汎用性と頑健性を示す。コードと量子化モデルはhttps://github.com/enyac-group/Quamba。

要約(オリジナル)

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 several 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.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL