重要提示

您正在查看 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 内存容量的模型（例如 NVGPT 5B 版本、LLaMA2 7B 版本或更大的模型）时，调整模型并行度至关重要。NeMo 结合了张量（层内）和流水线（层间）模型并行性。为了更深入地理解，请参阅“在 GPU 集群上使用 Megatron-LM 进行高效的大规模语言模型训练”(链接)。此工具可帮助修改模型并行度。在下载社区检查点并将其转换为 .nemo 格式后，如果用户希望进一步微调它，则此调整可能会变得至关重要。

了解检查点格式#

.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 模型类。

转换本地检查点#

只有最后一个检查点会自动以 .nemo 格式保存。如果需要中间训练检查点评估，则可能需要进行 .nemo 转换。为此，请参阅 <ADD convert_ckpt_to_nemo.py PATH> 处的脚本

python -m torch.distributed.launch --nproc_per_node=<tensor_model_parallel_size> * <pipeline_model_parallel_size> \
    examples/multimodal/convert_ckpt_to_nemo.py \
    --checkpoint_folder <path_to_PTL_checkpoints_folder> \
    --checkpoint_name <checkpoint_name> \
    --nemo_file_path <path_to_output_nemo_file> \
    --tensor_model_parallel_size <tensor_model_parallel_size> \
    --pipeline_model_parallel_size <pipeline_model_parallel_size>

转换社区检查点#

不支持将社区检查点转换为 NeMo ViT。

模型并行度调整#

ViT 检查点#

要将模型并行度从原始模型并行度大小调整为新的模型并行度大小（注意：NeMo ViT 目前仅支持 pipeline_model_parallel_size=1）

python examples/nlp/language_modeling/megatron_change_num_partitions.py \
 --model_file=/path/to/source.nemo \
 --target_file=/path/to/target.nemo \
 --tensor_model_parallel_size=??? \
 --target_tensor_model_parallel_size=??? \
 --pipeline_model_parallel_size=-1 \
 --target_pipeline_model_parallel_size=1 \
 --precision=32 \
 --model_class="nemo.collections.vision.models.megatron_vit_classification_models.MegatronVitClassificationModel" \
 --tp_conversion_only
在本页中