Quamba2: Un marco robusto y escalable de cuantización posentrenamiento para modelos de espacio de estados selectivos

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

Resumen

Los Modelos de Espacio de Estados (SSMs, por sus siglas en inglés) están surgiendo como una alternativa convincente a los Transformers debido a su uso constante de memoria y alto rendimiento. Sin embargo, escalar SSMs en servicios en la nube o dispositivos con recursos limitados es un desafío debido a sus requisitos de almacenamiento y potencia computacional. Para superar esto, la cuantización de SSMs con formatos de datos de bajo ancho de bits puede reducir el tamaño del modelo y beneficiarse de la aceleración por hardware. Dado que los SSMs son propensos a errores inducidos por la cuantización, esfuerzos recientes se han centrado en optimizar un modelo o ancho de bits específico para mejorar la eficiencia sin sacrificar el rendimiento. No obstante, configuraciones de ancho de bits distintas son esenciales para diferentes escenarios, como W4A8 para aumentar la velocidad de decodificación en lotes grandes, y W4A16 para mejorar la velocidad de generación en aplicaciones de prompts cortos para un solo usuario. Con este fin, presentamos Quamba2, compatible con W8A8, W4A8 y W4A16 para las arquitecturas Mamba1 y Mamba2, abordando la creciente demanda de implementación de SSMs en diversas plataformas. Basándonos en la preservación del orden de canales y la persistencia de activaciones de los SSMs, proponemos un enfoque offline para cuantizar las entradas de una recurrencia lineal en 8 bits mediante la clasificación y agrupación de la entrada x, combinado con una cuantización por grupo de estados para los parámetros dependientes de la entrada B y C. Para garantizar la invariancia computacional en la salida del SSM, reorganizamos los pesos offline según la secuencia de agrupación. Los experimentos muestran que Quamba2-8B supera a varios métodos de cuantización de SSMs de última generación y ofrece aceleraciones de 1.3 veces y 3 veces en las etapas de prellenado y generación, respectivamente, mientras reduce la memoria en 4 veces con solo una caída promedio del 1.6% en la precisión. La evaluación en MMLU demuestra la generalización y robustez de nuestro marco. El código y los modelos cuantizados se publicarán en: 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.