重要

您正在查看 NeMo 2.0 文档。此版本引入了 API 的重大更改和一个新的库，NeMo Run。我们目前正在将所有功能从 NeMo 1.0 移植到 2.0。有关以前版本或 2.0 中尚未提供的功能的文档，请参考 NeMo 24.07 文档。

Mamba 2#

状态空间模型 (SSM) 最近兴起成为 Transformer 的一种有前景的替代方案。SSM 具有诸如相对于序列长度的线性时间复杂度和用于推理的恒定缓存大小等优点。这些特性使得能够处理更长的序列和更高的吞吐量。尽管有这些优点，但仅 SSM 可能不如 Transformer 在需要强大的复制或上下文学习能力的任务上。

为了利用两种方法的优势，SSM 混合模型在其架构中结合了 MLP、Transformer 和 SSM 块。正如 NVIDIA 的一项研究 中强调的那样，由于包含 SSM 块，这些混合模型优于相同大小的传统 Transformer，从而实现了更快的推理时间。基于实验结果，Mamba2 混合模型不仅在性能上超越了 Transformer 基线，而且还受益于更高的计算效率。

《Transformers are SSMs》 论文中讨论的 Mamba2 模型有五种不同的尺寸可供选择：1.3 亿、3.7 亿、7.8 亿、13 亿和 27 亿参数。Mamba2 混合模型以及 NVIDIA 发布的 Mamba2 基线以 80 亿参数大小提供。

NeMo 2.0 预训练配方#

我们提供了预定义的配方，用于使用 NeMo 2.0 和 NeMo-Run 预训练以下尺寸的 Mamba2 和混合模型：1.3 亿、3.7 亿、7.8 亿、13 亿、27 亿、80 亿和 Hybrid-8B。这些配方为 nemo.collections.llm API 函数之一配置了一个 run.Partial，这些 API 函数是在 NeMo 2.0 中引入的。这些配方托管在 recipes <https://github.com/NVIDIA/NeMo/tree/main/nemo/collections/llm/recipes> 文件夹中（例如 mamba_130m.py）。

注意

预训练配方使用 MockDataModule 作为 data 参数。您需要将 MockDataModule 替换为您自定义的数据集。

我们在下面提供了一个示例，说明如何调用默认配方并覆盖数据参数

from nemo.collections import llm

pretrain = llm.mamba2_130m.pretrain_recipe(
    tokenizer_model="/path/to/tokenizer/model"
    name="mamba2_130m_pretraining",
    dir=f"/path/to/checkpoints",
    num_nodes=1,
    num_gpus_per_node=8,
)

dataloader = a_function_that_configures_your_custom_dataset(
    gbs=gbs,
    mbs=mbs,
    seq_length=pretrain.model.config.seq_length,
)
pretrain.data = dataloader

注意

对于 Mamba2 和混合模型，如果分词器在 Hugging Face 模型卡上不可用（对于 8B 和 Hybrid 8B 模型就是这种情况），您应该提供分词器模型的路径（默认为 None）。对于其他变体，请设置为 None）。分词器模型位于此处。

注意

配方中的配置是使用 NeMo-Run run.Config 和 run.Partial 配置对象完成的。请查看 NeMo-Run 文档 以了解有关其配置和执行系统的更多信息。

一旦您准备好最终配置，您就可以在任何 NeMo-Run 支持的执行器上执行它。最简单的是本地执行器，它只是在单独的进程中本地运行预训练。您可以按如下方式使用它

import nemo_run as run

run.run(pretrain, executor=run.LocalExecutor())

此外，您还可以按如下方式在同一个 Python 进程中直接运行它

run.run(pretrain, direct=True)

NeMo 2.0 微调配方#

与预训练类似，我们为上述 Mamba2 和混合模型提供了微调配方。微调配方托管在 recipes <https://github.com/NVIDIA/NeMo/tree/main/nemo/collections/llm/recipes> 文件夹中的相同位置（例如 mamba_130m.py）。

注意

微调配方使用 SquadDataModule（专为 SQUAD 数据集设计）作为 data 参数。您需要将 SquadDataModule 替换为您自定义的数据集。

我们在下面提供了一个示例，说明如何调用默认配方并覆盖数据参数

from nemo.collections import llm

pretrain = llm.mamba2_130m.finetune_recipe(
    resume_path="/path/to/nemo/checkpoint",
    tokenizer_model="/path/to/tokenizer/model",
    name="mamba2_130m_finetuning",
    dir=f"/path/to/checkpoints",
    num_nodes=1,
    num_gpus_per_node=8,
)

dataloader = a_function_that_configures_your_custom_dataset(
    gbs=gbs,
    mbs=mbs,
    seq_length=pretrain.model.config.seq_length,
)
pretrain.data = dataloader

注意

对于 Mamba2 和混合模型，您应该为所有模型提供您的 NeMo 检查点到 resume_path。此外，如果分词器在 Hugging Face 模型卡上不可用（8B 和 Hybrid 8B 模型就是这种情况），您应该将分词器模型路径提供给 tokenizer_model 参数（对于其他变体，请设置为 None）。分词器模型位于此处。

注意

配方中的配置是使用 NeMo-Run run.Config 和 run.Partial 配置对象完成的。请查看 NeMo-Run 文档 以了解有关其配置和执行系统的更多信息。

一旦您准备好最终配置，您就可以在任何 NeMo-Run 支持的执行器上执行它。最简单的是本地执行器，它只是在单独的进程中本地运行微调。您可以按如下方式使用它

import nemo_run as run

run.run(finetune, executor=run.LocalExecutor())

此外，您还可以按如下方式在同一个 Python 进程中直接运行它

run.run(finetune, direct=True)

下面提供了一个全面的预训练和微调配方列表，我们目前支持或计划很快支持，以供参考

配方

状态

Mamba2 130M

是

Mamba2 370M

是

Mamba2 780M

是

Mamba2 1.3B

是

Mamba2 2.7B

是

Mamba2 8B

是

Mamba2 Hybrid-8B

是
本页内容