通过组感知 SSM 剪枝实现高效的混合语言模型压缩

结合了注意力机制和状态空间模型 (SSM) 的混合 LLM 架构实现了最先进的精度和运行时性能。最近的研究表明，将压缩和蒸馏应用于仅注意力模型可以以较低的训练成本产生更小、更准确的模型。在这项工作中，我们探索了压缩混合架构的有效性。我们引入了一种新颖的组感知剪枝策略，该策略保留了 SSM 块的结构完整性及其序列建模能力。此外，我们证明了与传统方法相比，这种 SSM 剪枝对于实现更高的精度和推理速度是必要的。我们的压缩方案结合了 SSM、FFN、嵌入维度和层剪枝，然后进行基于知识蒸馏的再训练，类似于 MINITRON 技术。使用这种方法，我们将 Nemotron-H 8B 混合模型压缩到 40 亿参数，训练 token 减少了 40 倍。由此产生的模型超越了类似大小模型的精度，同时实现了快 2 倍的推理速度，显着推进了帕累托前沿。