Quamba2: Надежная и Масштабируемая Framework для Посттренировочной Квантования Селективных Моделей с Состояниями

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

March 28, 2025

Авторы: Hung-Yueh Chiang, Chi-Chih Chang, Natalia Frumkin, Kai-Chiang Wu, Mohamed S. Abdelfattah, Diana Marculescu

cs.AI

Аннотация

Модели пространства состояний (SSM) становятся привлекательной альтернативой трансформерам благодаря постоянному использованию памяти и высокой производительности. Однако масштабирование SSM на облачных сервисах или устройствах с ограниченными ресурсами затруднено из-за их требований к хранению и вычислительной мощности. Для решения этой проблемы квантование SSM с использованием форматов данных с низкой разрядностью может уменьшить размер модели и воспользоваться преимуществами аппаратного ускорения. Поскольку SSM подвержены ошибкам, вызванным квантованием, последние исследования сосредоточены на оптимизации конкретной модели или разрядности для повышения эффективности без ущерба для производительности. Тем не менее, различные конфигурации разрядности необходимы для разных сценариев, например, W4A8 для увеличения скорости декодирования при больших пакетах и W4A16 для повышения скорости генерации в приложениях с короткими запросами для одного пользователя. В связи с этим мы представляем Quamba2, совместимый с W8A8, W4A8 и W4A16 для архитектур Mamba1 и Mamba2, что отвечает растущему спросу на развертывание SSM на различных платформах. Основываясь на сохранении порядка каналов и устойчивости активаций в SSM, мы предлагаем оффлайн-подход для квантования входов линейной рекурсии в 8-битном формате путем сортировки и кластеризации для входного вектора x, в сочетании с квантованием по группам состояний для входно-зависимых параметров B и C. Чтобы обеспечить инвариантность вычислений на выходе SSM, мы переупорядочиваем веса оффлайн в соответствии с последовательностью кластеризации. Эксперименты показывают, что Quamba2-8B превосходит несколько современных методов квантования SSM и обеспечивает ускорение в 1.3 раза и 3 раза на этапах предварительного заполнения и генерации соответственно, при этом сокращая использование памяти в 4 раза с потерей точности всего на 1.6%. Оценка на MMLU демонстрирует универсальность и устойчивость нашего фреймворка. Код и квантованные модели будут доступны по адресу: 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.