重要提示

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

检查点#

在本节中，我们将介绍 NVIDIA NeMo 与检查点管理相关的四项关键功能

  1. 检查点加载：使用 restore_from() 方法加载本地 .nemo 检查点文件。

  2. 部分检查点转换：将部分训练的 .ckpt 检查点转换为 .nemo 格式。

  3. 社区检查点转换：将来自社区来源（如 HuggingFace）的检查点转换为 .nemo 格式。

  4. 模型并行度调整：修改模型并行度以有效地训练超出单个 GPU 内存的模型。NeMo 采用张量（层内）和流水线（层间）模型并行性。通过 “在 GPU 集群上使用 Megatron-LM 进行高效的大规模语言模型训练” 深入了解。此工具帮助调整模型并行度，以适应需要在更大的 GPU 阵列上部署以应对内存限制的用户。

理解检查点格式#

.nemo 检查点本质上是一个 tar 文件，它捆绑了模型配置（以 YAML 文件形式给出）、模型权重以及其他相关工件，如分词器模型或词汇表文件。这种整合的设计简化了共享、加载、调整、评估和推理。

相反，在 PyTorch Lightning 训练期间创建的 .ckpt 文件包含模型权重和优化器状态，通常用于从暂停处继续训练。

后续部分阐明了上述功能的说明，特别是为部署完全训练的检查点以进行评估或额外微调而定制的说明。

加载本地检查点#

默认情况下，NeMo 以 .nemo 格式保存训练模型的检查点。要在训练期间手动保存模型，请使用

model.save_to(<checkpoint_path>.nemo)

要加载本地 .nemo 检查点

import nemo.collections.multimodal as nemo_multimodal
model = nemo_multimodal.models.<MODEL_BASE_CLASS>.restore_from(restore_path="<path/to/checkpoint/file.nemo>")

将 <MODEL_BASE_CLASS> 替换为适当的 MM 模型类。

转换社区检查点#

CLIP 检查点#

要迁移社区检查点，请使用以下命令

torchrun --nproc-per-node=1 /opt/NeMo/scripts/checkpoint_converters/convert_clip_hf_to_nemo.py \
    --input_name_or_path=openai/clip-vit-large-patch14 \
    --output_path=openai_clip.nemo \
    --hparams_file=/opt/NeMo/examples/multimodal/vision_language_foundation/clip/conf/megatron_clip_VIT-L-14.yaml

确保 NeMo hparams 文件具有正确的模型架构参数，位于 path/to/saved.yaml。示例可以在 examples/multimodal/foundation/clip/conf/megatron_clip_config.yaml 中找到。

转换后，您可以使用以下命令验证模型

wget https://upload.wikimedia.org/wikipedia/commons/0/0f/1665_Girl_with_a_Pearl_Earring.jpg
torchrun --nproc-per-node=1 /opt/NeMo/examples/multimodal/vision_language_foundation/clip/megatron_clip_infer.py \
    model.restore_from_path=./openai_clip.nemo \
    image_path=./1665_Girl_with_a_Pearl_Earring.jpg \
    texts='["a dog", "a boy", "a girl"]'

它应该为 “a girl” 标签生成高概率。例如

Given image's CLIP text probability:  [('a dog', 0.0049710185), ('a boy', 0.002258187), ('a girl', 0.99277073)]
在本页中