使用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 examples。

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

训练愉快！

更新: