DeepSpeed通用检查点：实用指南

DeepSpeed通用检查点功能是用于以高效且灵活的方式保存和加载模型检查点的强大工具，能够实现跨不同模型架构、不同并行技术和训练配置的无缝模型训练延续和微调。本教程专为初学者和经验丰富的用户量身定制，提供有关如何在您的DeepSpeed驱动应用程序中利用通用检查点的分步指南。本教程将指导您完成创建ZeRO检查点、将其转换为通用格式以及使用这些通用检查点恢复训练的过程。这种方法对于利用预训练模型和促进跨不同设置的无缝模型训练至关重要。

通用检查点简介

DeepSpeed中的通用检查点抽象了保存和加载模型状态、优化器状态和训练调度程序状态的复杂性。此功能旨在开箱即用，只需最少的配置即可，支持各种规模和类型的模型，从小规模模型到大规模分布式模型，这些模型采用不同的并行拓扑结构，在多个GPU和其他加速器上进行训练。

先决条件

在开始之前，请确保您拥有以下内容

  • 已安装DeepSpeed，可以通过pip install deepspeed进行安装。
  • 一个使用DeepSpeed进行分布式训练的模型训练脚本。

如何使用DeepSpeed通用检查点

按照以下三个简单步骤操作

步骤1：创建ZeRO检查点

利用DeepSpeed通用检查点的第一步是创建一个ZeRO检查点。ZeRO（零冗余优化器）是DeepSpeed中的一项内存优化技术，它允许高效训练大型模型。要创建ZeRO检查点，您需要

  • 使用ZeRO优化器初始化您的模型与DeepSpeed。
  • 将您的模型训练到所需状态（迭代）。
  • 使用DeepSpeed的检查点功能保存检查点。

步骤2：将ZeRO检查点转换为通用格式

获得ZeRO检查点后，下一步是将其转换为通用格式。此格式旨在灵活且与不同的模型架构和DeepSpeed配置兼容。要转换检查点

  • 使用DeepSpeed提供的ds_to_universal.py脚本。
  • 指定ZeRO检查点的路径和通用检查点的所需输出路径。
python ds_to_universal.py --input_folder /path/to/zero/checkpoint --output_folder /path/to/universal/checkpoint

此脚本将处理ZeRO检查点并在通用格式中生成一个新的检查点。传递--help标志以查看其他选项。

步骤3：使用通用检查点恢复训练

准备好通用检查点后，您现在可以恢复训练，可能使用不同的并行拓扑或训练配置。为此，请将--universal-checkpoint添加到您的DeepSpeed配置（json）文件中

结论

DeepSpeed通用检查点简化了模型状态的管理，使保存、加载和跨不同训练会话和并行技术传输模型状态变得更加容易。通过遵循本教程中概述的步骤，您可以将通用检查点集成到您的DeepSpeed应用程序中，从而增强您的模型训练和开发工作流程。

有关更详细的示例和高级配置，请参阅Megatron-DeepSpeed示例。

有关DeepSpeed通用检查点的技术深入信息，请参阅arxiv手稿和博客。

训练愉快！

更新：