重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

长上下文配方#

长上下文模型训练增强了大型语言模型处理长上下文输入的能力。更长的序列长度有利于许多 NLP 任务，例如文档级摘要、长文档分类和长文档问答。NeMo 框架提供了配方来训练长上下文（16k 和 64k 序列长度）模型，如 Llama-3、Mixtral 和 Nemotron。上下文并行 (CP) 是 NeMo 框架中用于训练长上下文模型的关键技术。此技术允许训练和推理具有更长序列长度的模型。CP 使 NeMo 框架能够在其当前配方中支持 16k 和 64k 序列长度，并有可能扩展到更长的序列长度，例如 100 万。

上下文并行

NeMo 中提供的长上下文配方#

NeMo 2.0 提供了经过测试的配方来训练长上下文模型。这些配方在 NeMo LLM 配方目录中提供。目前，所有可用的配方都用于预训练，因为目前还没有微调用例。

Llama 3、Mixtral 和 Nemotron 模型的下表显示了不同序列长度以及 NeMo 中相应的配方。

Llama 3#

序列长度	8B	70B
16k	是	是
64k	是	是

Mixtral#

序列长度	8x7B
16k	是
64k	是

Nemotron 4#

序列长度	15B	22B
16k	是	是
64k	是	是

使用长上下文配方#

注意

对于长上下文配方，由于目前还没有实际的长上下文微调数据集，我们只支持预训练任务。

我们列出了两种访问长上下文配方的方法：使用 NeMo-Run API 和 NeMo-Run CLI。两者都包装了相同的 API 代码，但 NeMo-Run CLI 版本旨在易于使用。

通过 NeMo-Run API 使用#

注意

有关 NeMo-Run 的详细信息，请参阅 NeMo-Run 快速入门。以下是简明版本，重点介绍 NeMo 2.0 中长上下文配方的用法。

在 NeMo 2.0 中，我们包含了一个示例来演示我们现有的预训练配方。例如，您可以按如下方式使用 llama3_8b_16k 脚本

python pretraining.py --slurm --recipe llama3_8b_16k

--slurm 是非本地 Slurm 作业提交所必需的。此外，您必须在此示例 Python 文件中配置 Slurm 帐户设置。
模型类型和模型大小可以替换为 NeMo 2.0 中提供的其他配方（例如，模型大小为 8x7b_16k 的 mixtral）。

自定义参数#

要使用自定义参数，您可以轻松修改示例。例如，您可以在实验执行部分之前添加 pretrain.trainer.max_steps=2000 以覆盖此配方的最大步数。

通过 NeMo-Run CLI 使用#

注意

NeMo-Run CLI 目前仅在本地环境中工作。要在 Slurm 集群中使用它，请参阅 NeMo-Run 示例。

您可以通过 NeMo-Run CLI 使用这些配方。有关 NeMo-Run CLI 安装说明，请参阅 NeMo-Run CLI README。

例如

nemo llm pretrain --factory llama3_8b_16k

llama3_8b_16k 可以替换为 NeMo 2.0 中包含的其他配方（例如 mixtral_8x7b_16k）

注意

当使用多个进程（即在多个 GPU 上）启动配方时，请在命令中添加 -y 选项，以避免用户确认提示。例如，nemo llm pretrain --factory llama3_8b -y。

自定义参数#

您可以覆盖配方中的任何参数

nemo llm pretrain --factory llama3_8b_16k trainer.max_steps=2000

要从检查点继续训练，请将 resume.resume_from_path="to/some/path" 选项添加到命令中。您还可以通过这种方式指定其他选项，如序列长度。例如

nemo llm pretrain --factory llama3_8b_16k resume.resume_from_path="to/some/path" model.config.seq_length=2048 data.seq_length=2048

有关运行配方的更多详细信息，请参阅 Github 预训练 README。

添加新配方#

有关如何添加新配方的说明，请参阅 Github add-recipe README。

训练长上下文模型管线#

通常，训练长上下文模型涉及两个步骤：从头开始预训练，并通过持续预训练多次扩展先前模型的上下文长度。在 NeMo 2.0 中，可以使用 NeMo-Run CLI 通过以下步骤完成此过程。

通过 NeMo-Run API 使用#

从头开始预训练 Llama 3 模型#

以下是使用此脚本从头开始预训练序列长度为 2048 的 Llama 3 模型的示例。 llama3_8b 配方的默认序列长度为 8192。首先，您需要使用以下代码覆盖 llama3_8b 中使用的 seq_length

pretrain.model.config.seq_length=2048
pretrain.data.seq_length=2048

然后您可以使用以下命令启动作业

python pretraining.py --slurm --recipe llama3_8b

运行上述命令后，您将在目录中获得一个序列长度为 2048 的预训练模型检查点。

扩展预训练检查点的上下文长度#

要扩展模型的上下文长度，您可以从检查点继续训练。例如，您可以使用 16k 的序列长度

python pretraining.py --slurm --recipe llama3_8b_16k

接下来，您可以从 16k 检查点继续训练到 64k

python pretraining.py --slurm --recipe llama3_8b_64k

对于从不同目录保存和加载，请添加

pretrain.log.log_dir="your/second/log/ckpt/dir/here"
pretrain.resume.resume_from_path="your/first/log/ckpt/dir/here"

对于其他上下文长度，您需要为不同的并行配置添加相应的配方，或覆盖现有配方。有关详细信息，请参阅 GitHub add-recipe README。

通过 NeMo-Run CLI 使用#

注意

NeMo-Run CLI 目前仅在本地环境中工作。要在 Slurm 集群中使用，请参阅 NeMo-Run 示例。