混合Mamba-Transformer模型,推理速度提升3倍且更省内存。
Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models
- 用Mamba层替代Transformer大部分自注意力层,实现恒定计算与内存开销。
- 8B/56B模型在精度上优于或持平主流开源模型,推理快3倍。
- 通过剪枝与蒸馏压缩出47B基座模型,速度提升20%仍保高精度。
随着推理时扩展对增强推理能力的重要性日益凸显,构建高效推理的模型变得愈发关键。本文提出Nemotron-H,一组8B和56B/47B混合Mamba-Transformer模型,旨在降低特定精度下的推理成本。通过将通用Transformer架构中的多数自注意力层替换为执行恒定计算、每生成一个词元仅需恒定内存的Mamba层,实现这一目标。实验表明,Nemotron-H模型在精度上优于或持平其他同规模的先进开源Transformer模型(如Qwen-2.5-7B/72B和Llama-3.1-8B/70B),同时推理速度最高提升3倍。为进一步提升推理速度并减少推理时内存占用,我们基于56B模型,采用一种名为MiniPuzzle的新压缩技术(结合剪枝与蒸馏)训练出Nemotron-H-47B-Base模型。该模型在精度上与56B模型相当,但推理速度提升20%。此外,我们提出基于FP8的训练方案,结果显示其性能可与BF16训练相当,并用于训练56B模型。所有Nemotron-H基座模型检查点已在Hugging Face和NeMo中发布。
原文摘要 · Abstract (English)
As inference-time scaling becomes critical for enhanced reasoning capabilities, it is increasingly becoming important to build models that are efficient to infer. We introduce Nemotron-H, a family of 8B and 56B/47B hybrid Mamba-Transformer models designed to reduce inference cost for a given accuracy level. To achieve this goal, we replace the majority of self-attention layers in the common Transformer model architecture with Mamba layers that perform constant computation and require constant memory per generated token. We show that Nemotron-H models offer either better or on-par accuracy compared to other similarly-sized state-of-the-art open-sourced Transformer models (e.g., Qwen-2.5-7B/72B and Llama-3.1-8B/70B), while being up to 3$\times$ faster at inference. To further increase inference speed and reduce the memory required at inference time, we created Nemotron-H-47B-Base from the 56B model using a new compression via pruning and distillation technique called MiniPuzzle. Nemotron-H-47B-Base achieves similar accuracy to the 56B model, but is 20% faster to infer. In addition, we introduce an FP8-based training recipe and show that it can achieve on par results with BF16-based training. This recipe is used to train the 56B model. We are releasing Nemotron-H base model checkpoints with support in Hugging Face and NeMo.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。