Compression efficace de modèles de langage hybrides par élagage SSM basé sur les groupes

Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

April 15, 2025

Auteurs: Ali Taghibakhshi, Sharath Turuvekere Sreenivas, Saurav Muralidharan, Marcin Chochowski, Yashaswi Karnati, Raviraj Joshi, Ameya Sunil Mahabaleshwarkar, Zijia Chen, Yoshi Suhara, Oluwatobi Olabiyi, Daniel Korzekwa, Mostofa Patwary, Mohammad Shoeybi, Jan Kautz, Bryan Catanzaro, Ashwath Aithal, Nima Tajbakhsh, Pavlo Molchanov

cs.AI

Résumé

Les architectures hybrides de grands modèles de langage (LLM) qui combinent les mécanismes d'Attention et les modèles d'espace d'état (SSM) atteignent des performances de pointe en termes de précision et de temps d'exécution. Des travaux récents ont montré que l'application de techniques de compression et de distillation aux modèles basés uniquement sur l'Attention permet d'obtenir des modèles plus petits et plus précis à une fraction du coût d'entraînement. Dans ce travail, nous explorons l'efficacité de la compression des architectures hybrides. Nous introduisons une nouvelle stratégie d'élagage sensible aux groupes qui préserve l'intégrité structurelle des blocs SSM et leurs capacités de modélisation de séquences. De plus, nous démontrons la nécessité d'un tel élagage des SSM pour obtenir une amélioration de la précision et de la vitesse d'inférence par rapport aux approches traditionnelles. Notre méthode de compression combine l'élagage des SSM, des couches FFN, des dimensions d'embedding et des couches, suivi d'un réentraînement basé sur la distillation de connaissances, similaire à la technique MINITRON. En utilisant cette approche, nous compressons le modèle hybride Nemotron-H 8B jusqu'à 4 milliards de paramètres avec jusqu'à 40 fois moins de tokens d'entraînement. Le modèle résultant surpasse la précision des modèles de taille similaire tout en atteignant une inférence 2 fois plus rapide, repoussant significativement la frontière de Pareto.

English

Hybrid LLM architectures that combine Attention and State Space Models (SSMs) achieve state-of-the-art accuracy and runtime performance. Recent work has demonstrated that applying compression and distillation to Attention-only models yields smaller, more accurate models at a fraction of the training cost. In this work, we explore the effectiveness of compressing Hybrid architectures. We introduce a novel group-aware pruning strategy that preserves the structural integrity of SSM blocks and their sequence modeling capabilities. Furthermore, we demonstrate the necessity of such SSM pruning to achieve improved accuracy and inference speed compared to traditional approaches. Our compression recipe combines SSM, FFN, embedding dimension, and layer pruning, followed by knowledge distillation-based retraining, similar to the MINITRON technique. Using this approach, we compress the Nemotron-H 8B Hybrid model down to 4B parameters with up to 40x fewer training tokens. The resulting model surpasses the accuracy of similarly-sized models while achieving 2x faster inference, significantly advancing the Pareto frontier.