Snapckpt加速包

SnapCheckpoint（Snapckpt）是由息壤一体化智算服务平台 - 一站式智算服务平台提供的针对大模型训练场景提供的高性能checkpoint框架。Snapckpt提供了高效便捷的断点存储、加载机制。

1、Snapckpt介绍

Snapckpt是一种为大模型训练打造的易用、可拓展、高性能的断点解决方案。Snapckpt提供了高效便捷的断点存储、加载机制，其主要有以下几大优势：

1. 异步断点保存：通过优化断点保存流程，使得断点保存过程异步化，尽可能减小断点保存过程对于训练的中断，减少训练阶段耗时。

2. 分布式断点存储：对模型机优化器参数进行分布式存储，避免集中式存储带来的额外通信开销，大幅降低存储与加载断点的耗时。

3. 多框架支持：支持Megatron-LM、Deepspeed两大主流大模型训练框架。

4. 简单易用：安装及使用方式简单，用户仅需简单步骤即可使用，提升使用效率。

2、背景信息

在大规模分布式模型训练过程中，系统可能因硬件故障或软件异常导致训练中断。为确保训练进度可恢复，业界普遍采用周期性保存检查点(Checkpoint)的方案。值得注意的是，检查点操作耗时与模型参数量呈正相关关系，当面对参数量达百亿甚至千亿级别的大模型时，每次检查点保存往往需要耗费数分钟至十余分钟不等。特别是在使用Megatron-LM框架或原生PyTorch进行训练时，检查点保存过程会强制暂停训练任务，导致宝贵的计算资源被闲置。因此，开发高效的检查点机制以降低时间成本和资源浪费，已成为当前大规模模型训练亟待解决的关键问题。

Snapckpt采用多阶段异步断点存储机制，降低断点存储耗时，减少训练中断带来的影响，从而提升训练速度，提升计算资源有效利用率。

3、使用方法

3.1、安装Snapckpt

1、Snapckpt Python包下载地址：https://shanghai15.zos.ctyun.cn/zhisuan-sre/huiju-data/Ctyun%20Snapckpt.zip

1. 下载后进入到工程根目录

2. 安装命令：

python3 setup.py build && python3 setup.py install

至此，Snapckpt的软件包安装完成，可以开始使用（具体使用方法见“Snapckpt使用方法”章节）。

3.2、Snapckpt使用方法

3.2.1、Megatron模式使用方法

在使用Megatron-LM训练框架训练模型时，您可以使用 snapckpt.megatron_dist_ckpt.save_checkpoint 和 snapckpt.megatron_dist_ckpt.load_checkpoint 替换原生Megatron-LM的 megatron.checkpointing.save_checkpoint 和 megatron.checkpointing.load_checkpoint。

# from megatron.checkpointing import save_checkpoint
from snapckpt.megatron_dist_ckpt import save_checkpoint

# from megatron.checkpointing import load_checkpoint
from snapckpt.megatron_dist_ckpt import load_checkpoint

3.2.2、Deepspeed模式使用方法

首先，需要初始化Snapckpt的DeepSpeedCheckpointer，初始化方式如下：

from snapckpt.deepspeed import (
    DeepSpeedCheckpointer,
    StorageType,
)
checkpointer = DeepSpeedCheckpointer(model, checkpoint_dir)

初始化后，您可以使用 checkpointer.save_checkpoint 和 checkpointer.load_checkpoint 替换原生Deepspeed的 save_model 和 load_checkpoint。

# save_model(checkpoint_dir)
checkpointer.save_checkpoint(checkpoint_dir, tag=step, storage_type=StorageType.MEMORY)
checkpointer.save_checkpoint(checkpoint_dir, tag=step, storage_type=StorageType.DISK)

# load_checkpoint(model, checkpoint_dir)
checkpointer.load_checkpoint(checkpoint_dir)

4、性能参考

以Deepseek R1 Distill Qwen 14B模型，单机八卡运行环境为例，相较于同步断点保存而言，Snapckpt的性能提升测试参考如下：