Quamba2: Een robuust en schaalbaar raamwerk voor post-training kwantisering voor selectieve toestandsruimtemodellen

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

March 28, 2025

Auteurs: Hung-Yueh Chiang, Chi-Chih Chang, Natalia Frumkin, Kai-Chiang Wu, Mohamed S. Abdelfattah, Diana Marculescu

cs.AI

Samenvatting

State Space Models (SSMs) doen zich steeds meer gelden als een overtuigend alternatief voor Transformers vanwege hun consistente geheugengebruik en hoge prestaties. Desondanks is het opschalen van SSMs op clouddiensten of apparaten met beperkte middelen een uitdaging vanwege hun opslagvereisten en rekenkracht. Om dit te overwinnen, kan het kwantiseren van SSMs met dataformaten met een laag bitbreedte het modelformaat verkleinen en profiteren van hardwareversnelling. Omdat SSMs gevoelig zijn voor kwantisatiefouten, hebben recente inspanningen zich gericht op het optimaliseren van een specifiek model of bitbreedte voor efficiëntie zonder in te leveren op prestaties. Echter, verschillende bitbreedteconfiguraties zijn essentieel voor verschillende scenario's, zoals W4A8 voor het verhogen van de decodersnelheid bij grote batches, en W4A16 voor het verbeteren van de generatiesnelheid in toepassingen met korte prompts voor een enkele gebruiker. Daarom presenteren we Quamba2, compatibel met W8A8, W4A8 en W4A16 voor zowel Mamba1- als Mamba2-backbones, waarmee we inspelen op de groeiende vraag naar SSM-implementatie op diverse platforms. Gebaseerd op de kanaalvolgorde-behoudende en activatie-persistente eigenschappen van SSMs, stellen we een offline aanpak voor om de invoer van een lineaire recurrentie te kwantiseren in 8-bit door sorteren en clusteren voor invoer x, gecombineerd met een per-staat-groep-kwantisatie voor invoerafhankelijke parameters B en C. Om compute-invariantie in de SSM-uitvoer te garanderen, herschikken we de gewichten offline volgens de clusteringvolgorde. De experimenten tonen aan dat Quamba2-8B verschillende state-of-the-art SSM-kwantisatiemethoden overtreft en respectievelijk 1,3x en 3x snelheidsverbeteringen biedt in de pre-filling- en generatiefasen, terwijl het een 4x geheugenreductie biedt met slechts een gemiddelde nauwkeurigheidsdaling van 1,6%. De evaluatie op MMLU toont de generaliseerbaarheid en robuustheid van ons framework. De code en gekwantiseerde modellen zullen worden vrijgegeven op: https://github.com/enyac-group/Quamba.

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 several state-of-the-art SSM quantization methods and delivers 1.3times and 3times speed-ups in the pre-filling and generation stages, respectively, while offering 4times 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.