NVIDIA NeMo 通过混合状态空间模型集成增强LLM能力 - Blockchain.News

NVIDIA NeMo 通过混合状态空间模型集成增强LLM能力

realtime news Jul 18, 2024 02:50

NVIDIA NeMo 引入了对混合状态空间模型的支持,显著提升了大语言模型的效率和能力。

NVIDIA NeMo 通过混合状态空间模型集成增强LLM能力

根据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