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

要約

状態空間モデル（SSM）は、一貫したメモリの使用と高性能のために、トランスの魅力的な代替品として浮上しています。
それにもかかわらず、クラウドサービスまたはリミテッドリソースデバイスでSSMをスケールアップすることは、そのストレージ要件と計算能力のために困難です。
これを克服するために、ビット幅データ形式が低いSSMを量子化すると、モデルのサイズを削減し、ハードウェアアクセラレーションから利益を得ることができます。
SSMは量子化によるエラーが発生しやすいため、最近の努力は、パフォーマンスを犠牲にすることなく、特定のモデルまたは効率のためにビット幅を最適化することに焦点を当てています。
ただし、大型バッチデコード速度を高めるためのW4A8や、単一のユーザーの短いプロンプトアプリケーションで生成速度を向上させるためのW4A16など、異なるシナリオには異なるビット幅構成が不可欠です。
この目的のために、Mamba1とMamba2の両方の骨格についてW8A8​​、W4A8、およびW4A16と互換性のあるQuamba2を提示し、さまざまなプラットフォームでのSSM展開の需要の増加に対処します。
SSMのチャネル順序の保存と活性化の持続性に基づいて、入力$ x $を選別およびクラスタリングすることにより、8ビットの線形再発の入力を量子化するオフラインアプローチを提案します。
SSM出力での計算不変性を確保するために、クラスタリングシーケンスに従って重みをオフラインで再配置します。
この実験では、Quamba2-8Bが2つの最先端の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 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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google