Efficient Hybrid Language Model Compression through Group-Aware SSM Pruning

Bryan Catanzaro (NVIDIA) · Mostofa Patwary (NVIDIA) · Mohammad Shoeybi (NVIDIA) · Pavlo Molchanov (NVIDIA) · Jan Kautz (NVIDIA) · ZIJIA CHEN (NVIDIA) · Yoshi Suhara (NVIDIA) · Ali Taghibakhshi (NVIDIA) · Sharath Turuvekere Sreenivas (NVIDIA) · Saurav Muralidharan (NVIDIA) · Marcin Chochowski (NVIDIA) · Yashaswi Karnati (NVIDIA) · Raviraj Joshi (Department of Computer Science, Indian Institute of Technology, Madras, Indian Institute of Technology, Madras) · Ameya Mahabaleshwarkar (NVIDIA) · Oluwatobi Olabiyi (NVIDIA) · Daniel Korzekwa (Nvidia) · Ashwath Aithal (NVIDIA) · Nima Tajbakhsh (Illinois Institute of Technology)
attentioncompressing architecturesdistillationembedding dimensionffngroup-aware pruninghybrid language modelsinference throughputknowledge distillationlayer pruningmamba layerspruningsequence modelingstate space modelsstructural integrityunified compression recipe

Hybrid language models that combine Attention and State Space Models (SSMs) have been shown to achieve state-of-the-art accuracy and runtime performance. Recent work has also demonstrated that applying pruning 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. To this end, we introduce a novel group-aware pruning method for Mamba layers that preserves the structural integrity of SSM blocks and their sequence modeling capabilities. We combine this method with FFN, embedding dimension, and layer pruning, along with knowledge distillation-based retraining to obtain a unified compression recipe for hybrid models. Using this recipe, we compress the Nemotron-H 8B Hybrid model down to 4B parameters with up to $40\times$ fewer training tokens compared to similarly-sized models. The resulting model surpasses the accuracy of similarly-sized models while achieving $\sim2\times$ faster inference throughput, significantly advancing the Pareto frontier.