Quamba2: Um Framework Robusto e Escalável de Quantização Pós-Treinamento para Modelos de Espaço de Estados Seletivos

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

March 28, 2025

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

cs.AI

Resumo

Os Modelos de Espaço de Estados (SSMs) estão surgindo como uma alternativa atraente aos Transformers devido ao uso consistente de memória e ao alto desempenho. No entanto, escalar SSMs em serviços de nuvem ou dispositivos com recursos limitados é desafiador devido aos requisitos de armazenamento e poder computacional. Para superar isso, a quantização de SSMs com formatos de dados de baixa largura de bits pode reduzir o tamanho do modelo e se beneficiar da aceleração de hardware. Como os SSMs são propensos a erros induzidos pela quantização, esforços recentes têm se concentrado em otimizar um modelo ou largura de bits específica para eficiência sem sacrificar o desempenho. No entanto, configurações distintas de largura de bits são essenciais para diferentes cenários, como W4A8 para aumentar a velocidade de decodificação em grandes lotes, e W4A16 para melhorar a velocidade de geração em aplicações de prompt curto para um único usuário. Para isso, apresentamos o Quamba2, compatível com W8A8, W4A8 e W4A16 para as arquiteturas Mamba1 e Mamba2, atendendo à crescente demanda de implantação de SSMs em várias plataformas. Com base na preservação da ordem dos canais e na persistência da ativação dos SSMs, propomos uma abordagem offline para quantizar as entradas de uma recorrência linear em 8 bits, ordenando e agrupando a entrada x, combinada com uma quantização por grupo de estados para os parâmetros dependentes da entrada B e C. Para garantir a invariância computacional na saída do SSM, reorganizamos os pesos offline de acordo com a sequência de agrupamento. Os experimentos mostram que o Quamba2-8B supera vários métodos de quantização de SSMs de última geração e oferece acelerações de 1,3 vezes e 3 vezes nas etapas de pré-preenchimento e geração, respectivamente, enquanto proporciona uma redução de memória de 4 vezes com apenas uma queda média de precisão de 1,6%. A avaliação no MMLU demonstra a generalizabilidade e robustez do nosso framework. O código e os modelos quantizados serão disponibilizados em: 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.