重要提示

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

将检查点配置从 NeMo 1.0 迁移到 NeMo 2.0#

NeMo 1.0（先前版本）#

在 NeMo 1.0 中，分布式检查点在 YAML 配置文件中配置。

# Distributed checkpoint setup
dist_ckpt_format: 'zarr' # Set to 'torch_dist' to use PyTorch distributed checkpoint format.
dist_ckpt_load_on_device: True # whether to load checkpoint weights directly on GPU or to CPU
dist_ckpt_parallel_save: False # if true, each worker will write its own part of the dist checkpoint
dist_ckpt_parallel_load: False # if true, each worker will load part of the dist checkpoint and exchange with NCCL. Might use some extra GPU memory
dist_ckpt_torch_dist_multiproc: 2 # number of extra processes per rank used during ckpt save with PyTorch distributed format
dist_ckpt_assume_constant_structure: False # set to True only if the state dict structure doesn't change within a single job. Allows caching some computation across checkpoint saves.
dist_ckpt_parallel_dist_opt: True # parallel save/load of a DistributedOptimizer. 'True' allows performant save and reshardable checkpoints. Set to 'False' only in order to minimize the number of checkpoint files.

NeMo 2.0（新版本）#

在 NeMo 2.0 中，这些设置从 MegatronStrategy 控制。

from nemo.collections import llm
from nemo import lightning as nl

strategy = nl.MegatronStrategy(
   save_ckpt_format='zarr',
   ckpt_load_on_device=True,
   ckpt_parallel_save=False,
   ckpt_parallel_load=False,
   ckpt_assume_constant_structure=False,
   ckpt_parallel_save_optim=False,
)

nl.Trainer(
   strategy=strategy,
   ...
)

迁移分布式检查点设置#

  1. 在您的 NeMo 1.0 YAML 配置文件中找到分布式检查点设置部分。

  2. 将 distributed checkpoint setup 设置传递到 MegatronStrategy

    strategy = nl.MegatronStrategy(
    save_ckpt_format='zarr',
    ckpt_load_on_device=True,
    ckpt_parallel_save=False,
    ckpt_parallel_load=False,
    ckpt_torch_dist_multiproc=2,
    ckpt_assume_constant_structure=False,
    ckpt_parallel_save_optim=False,
)

注意

NeMo 2.0 不支持非分布式检查点。

将 NeMo 1.0 检查点转换为 NeMo 2.0 检查点#

我们提供了一个脚本来将 NeMo 1.0 检查点转换为 NeMo 2.0 检查点。该脚本位于 scripts/checkpoint_converters/convert_nemo1_to_nemo2.py。

  1. NeMo 1.0 检查点采用 model_name.nemo tarball 格式。解压 tarball 后，您将看到以下结构

    model_name.nemo
├── model_config.yaml
├── model_weights
│   ├── distributed checkpointing directories/files in zarr or torch_dist format
│   ├── metadata.json
│   └── common.pt

  2. NeMo 2.0 检查点是一个具有以下结构的目录。context 目录中的内容可能是可选的

    model_name
├── context
│   ├── model_config.yaml
│   ├── io.json
│   └── tokenizer
├── weights
│   ├── distributed checkpointing directories/files in zarr or torch_dist format
│   ├── metadata.json
│   └── common.pt

  3. 转换脚本需要 NeMo 1.0 权重、NeMo 1.0 模型 ID（指定模型结构和配置）和 NeMo 1.0 分词器信息（sentencepiece tokenizer.model 或 Hugging Face 分词器 ID）才能将检查点转换为 NeMo 2.0 格式。该脚本将创建一个具有 NeMo 2.0 检查点结构的新目录。

    转换脚本使用 CPU 和 CPU 内存进行检查点转换。当您的 NeMo 1.0 检查点使用 Hugging Face 分词器时，转换脚本将从 Hugging Face 下载分词器。如果分词器来自 gated repo，您需要先登录 Hugging Face

    huggingface-cli login

    目前，仅支持 sentencepice 和 Hugging Face 分词器。

    以下命令应在 NeMo 容器内使用。

    • 您可以将包含权重和分词器信息的 model_name.nemo tarball 传递给脚本。我们以 meta-llama/Meta-Llama-3-8B 为例

    python /opt/NeMo/scripts/checkpoint_converters/convert_nemo1_to_nemo2.py \
--input_path=Meta-Llama-3-8B.nemo \
--output_path=your_output_dir \
--model_id=meta-llama/Meta-Llama-3-8B

    • 如果您有一个模型权重目录（其结构类似于 NeMo 1.0 检查点中的 model_weights 目录），则可以将权重目录传递给脚本。在这种情况下，脚本还需要分词器信息，因为权重目录不包含此信息。我们以 nemotron-3-8b-base-4k` 为例

    python /opt/NeMo/scripts/checkpoint_converters/convert_nemo1_to_nemo2.py \
--input_path=nemotron3-8b-extracted/model_weights \
--tokenizer_path=path_to_your_tokenizer_model.model \
--tokenizer_library=sentencepiece \
--output_path=your_output_dir \
--model_id=nvidia/nemotron-3-8b-base-4k

  4. 支持的模型：目前，我们已经验证了以下模型

    • meta-llama/Meta-Llama-3-8B

    • mistralai/Mixtral-8x7B-v0.1

    • nvidia/nemotron-3-8b-base-4k

    相同系列/结构和不同大小的模型应该可以工作，但尚未经过验证。模型转换仅适用于 NeMo 2.0 支持的模型。我们将在未来添加更多模型的验证。可在脚本 scripts/checkpoint_converters/convert_nemo1_to_nemo2.py 中找到可用模型 ID 列表。--model_id 参数应为以下之一

    • meta-llama/Llama-2-7b-hf

    • meta-llama/Llama-2-13b-hf

    • meta-llama/Llama-2-70b-hf

    • meta-llama/Meta-Llama-3-8B

    • meta-llama/Meta-Llama-3-70B

    • mistralai/Mixtral-8x7B-v0.1

    • mistralai/Mixtral-8x22B-v0.1

    • mistralai/Mistral-7B-v0.1

    • nvidia/nemotron-3-8b-base-4k

    • nemotron4-22b

    • nemotron4-15b

    • nemotron4-340b

在本页中