NVIDIA NeMo 通过混合状态空间模型集成增强LLM能力
realtime news Jul 18, 2024 02:50
NVIDIA NeMo 引入了对混合状态空间模型的支持，显著提升了大语言模型的效率和能力。
根据NVIDIA技术博客的报道，NVIDIA宣布在其NeMo框架中集成混合状态空间模型（SSM），这是人工智能领域的重要举措。此项发展有望提升大语言模型（LLM）的效率和能力。
基于Transformer模型的进展
自2017年引入transformer模型架构以来，AI计算性能迅速发展，使得更大更强的LLM得以创建。这些模型在智能聊天机器人、计算机代码生成甚至芯片设计中找到了应用。
为了支持这些先进LLM的训练，NVIDIA NeMo提供了一个用于构建、定制和部署LLM的端到端平台。NeMo内集成了Megatron-Core，这是一个基于PyTorch的库，提供了训练LLM所需的基本组件和优化。
状态空间模型的引入
NVIDIA的最新公告包括对状态空间模型（SSM）预训练和微调的支持。此外，NeMo现在还支持基于Google DeepMind描述的Griffin架构的模型训练。
替代模型架构的优势
尽管transformer模型通过注意机制在捕捉长距离依赖方面表现出色，但其计算复杂度随着序列长度的增加呈二次方增长，导致训练时间和成本增加。然而，SSM通过克服与注意机制相关的多种限制，提供了一种具有吸引力的替代方案。
SSM以其计算和内存方面的线性复杂度著称，使其在建模长距离依赖关系上更加高效。它们还提供了与transformer模型相当的高质量和准确性，并且在推理过程中需要更少的内存。
SSM在长序列训练中的效率
由于其高效处理序列建模任务的特点，SSM在深度学习社区中越来越受欢迎。例如，当序列长度增加到256K时，SSM的一个变体Mamba-2层比transformer层快18倍。
Mamba-2采用了结构化状态空间对偶（SSD）层，将SSM计算重新表述为矩阵乘法，从而利用NVIDIA Tensor Cores的性能。这使得Mamba-2在保持与transformer相竞争的质量和准确性的同时，训练速度更快。
增强性能的混合模型
通过结合SSM、SSD、RNN和transformer的混合模型可以利用每种架构的优点，并减轻各自的缺点。NVIDIA研究人员最近的一篇论文描述了混合Mamba-Transformer模型，这些模型在标准任务上的表现超过了纯transformer模型，预计推理时可快8倍。
这些混合模型还显示出更高的计算效率。随着序列长度的增加，训练混合模型所需的计算量的增长率远低于纯transformer模型。
未来展望
NVIDIA NeMo对SSM和混合模型的支持标志着实现新AI智能水平的重要一步。初期功能包括对Mamba-2、Griffin架构、混合模型组合以及各种模型微调的支持。未来的版本预计将包括更多的模型架构、性能优化并支持FP8训练。
欲了解更多详细信息，请访问NVIDIA技术博客。
Image source: Shutterstock