当今的大型语言模型（LLM）基于 2017 年推出的 Transformer 模型架构。自那时以来，AI 计算性能的快速进步使创建更大的基于 Transformer 的 LLM 成为可能，这极大地提高了其功能。基于 Transformer 的高级 LLM 正在为许多令人兴奋的应用提供支持，如智能聊天机器人、计算机代码生成和甚至 芯片设计 等。
训练先进的 LLM 需要高效且通用的软件堆栈。为此，NVIDIA NeMo 提供了一个端到端平台，用于构建、自定义和部署 LLM。Integrated 深入集成到 NeMo 框架中的是 Megatron-Core，一个基于 PyTorch 的库，提供了大规模训练 LLM 所需的基本组件和优化。随着模型开发者探索新的模型架构，NVIDIA 平台不断扩展，以实现创新。
今天，NVIDIA 宣布 NeMo 和 Megatron-Core 现在分别支持状态空间模型（SSM）的预训练和微调。另外，NeMo 现在还支持基于 Google DeepMind 所述的 Griffin 架构的模型训练。
为何要探索替代模型架构？
Transformer 模型擅长捕捉远程依赖项，因为它们通过著名的注意力机制实现了全局上下文理解，从而非常适合处理需要这种理解的任务。
但是，注意力的计算复杂性会随序列长度二次扩展，从而导致训练时间和训练成本的大幅增加。随着序列长度的增加，这些成本也会增加。此外，在推理期间，注意力需要存储键值对缓存（称为 KV 缓存），这些缓存会随着序列长度线性增长。这会导致内存占用随着序列长度的增加而增加。
最近，SSM 模型架构克服了一些注意力限制，成为处理序列建模任务的极具吸引力的模型架构。
SSM 可以实现长序列长度训练的更高效率
SSM 是一类模型，在深度学习社区中越来越受欢迎，作为基于注意力的 Transformer 模型的高效替代方案，适用于序列建模任务。
SSM 具有以下令人信服的特性：
- 线性复杂性：SSM 在计算和内存复杂性方面都是线性的，而注意力在这两方面都是二次的。这意味着，相比注意力，SSM 可以更高效地对序列中的长程依赖项建模。
- 高质量和高精度：与注意力一样，SSM 查看输入序列的标记，使模型能够专注于最相关的部分，从而实现与基于 Transformer 的模型相当的质量和准确性。
- 高效推理：SSM 只需存储常量大小的向量，而不是 KV 缓存，这使得推理变得更加内存高效，特别是在序列长度较长的情况下。
为说明 SSM 为更长序列长度提供的优势，下图显示了 Mamba-2 层 (稍后将在本文中介绍的状态空间模型变体) 与训练 Transformer 层相比的相对速度随着序列长度的增加。随着序列长度增加到 256K，Mamba-2 层的速度比 Transformer 层快 18 倍。
Transformer 模型维度为 4096，拥有 32 个头。Mamba-2 模型维度为 4096，状态维度为 128，分为 8 组
一些 SSM 变体在 AI 社区中已经流行起来，其中包括 Hyena、Mamba-1 和最近推出的 Mamba-2。
结构化状态空间对偶性和 Mamba-2
Mamba-2 作为最新版本脱颖而出，在多个基准测试中实现了非常高的准确性。其核心是一个新的结构化状态空间对偶（SSD）层。在实践中，这个 SSD 层重新表述了 Mamba-1 模型中使用的 SSM 数学运算。这种重新表述将 SSM 计算重构为矩阵乘法，从而使它们能够利用 NVIDIA Tensor Core 的重要矩阵乘法性能。
因此，与 Mamba-1 相比，Mamba-2 的训练速度要快得多。另外，在语言建模任务中，Mamba-2 还提供了与 Transformer 相比具有竞争力的质量和准确性。当在混合模型中将几个注意力层与 SSD 层相结合时，Mamba-2 可以产生更好的结果。
但是，纯粹的 SSM 并非没有限制。例如，它们在“干草堆”类型场景中很困难，这些场景需要在非常长的序列中精确地调用信息。
混合模型可以改善结果，同时提高性能
结合 SSM、SSD、RNN 和 Transformer 的混合模型可以充分发挥每个模型架构的优势，同时减轻其各自的弱点。
在最近的一篇论文中，包括 NVIDIA 应用深度学习研究 (ADLR) 团队成员在内的研究人员描述了混合 Mamba-Transformer 模型。在这些混合模型中，标准 Transformer 层和新型 SSM 层可以在任意配置中交织在一起。例如，本文中描述的 8B 混合模型有 56 层，其中包括 4 层自注意力层、24 层 Mamba-2 层和 28 层多层感知器 (MLP) 层。这些层的分配是这样的：首先出现 Mamba-2 层，然后是注意力层，最后 MLP 层均匀分布在整个模型中。
根据该论文，该团队评估的混合 8B Mamba-2-Hybrid 模型在所有 12 项标准任务中都超过了 8B Transformer。此外，该模型还“预计在推理时生成令牌的速度将提高 8 倍”。
除了在推理过程中提高执行任务的能力和显著的性能优势外，Mamba-2-Hybrid 模型还显示了更高的计算效率。在序列长度增加时，训练 8B Mamba-2-Hybrid 模型和训练 8B Transformer 模型所需的计算量进行了比较，结果如下图所示。
当序列长度为 2048 个令牌时，两者的计算需求大致相同，其中混合模型略有优势。然而，当序列长度扩展到多达 32768 个令牌时，8B Transformer 模型的计算需求将翻倍，而混合模型仅增长 13%。由于现代语言模型支持 1M 令牌及以上的序列长度，因此 SSM-Transformer 混合模型的这种优势只会增长。
支持新型模型架构的第一步
模型架构创新对于实现更高水平的智能至关重要。除了为构建基于 Transformer 的模型提供出色支持外，NeMo 和 Megatron-Core 现在还为社区提供了训练 Self-Supervised Models（SSMs）和 Single-Stage Detectors（SSDs）的能力，以及将其优势与 Transformer 模型优势相结合的混合模型的能力。
在此版本的 NeMo 中，我们提供了以下初始功能，以便社区能够快速开始试验：
- 支持 SSD 模型，包括Mamba-2。
- 支持 RG-LRU (Griffin 架构)
- 支持 Transformer/SSM 混合模型的组合。
- 微调支持 Recurrent Gemma（Griffin）模型、纯 Mamba-2 模型和 8B Mamba-2-Hybrid 模型。
- 分片和模型并行支持。
在即将发布的版本中，我们计划支持其他子二次模型架构，以及其他性能优化和 FP8 训练的支持。
