Snapckpt加速包
SnapCheckpoint（Snapckpt）是由息壤一體化智算服務平臺 - 一站式智算服務平臺提供的針對大模型訓練場景提供的高性能checkpoint框架。Snapckpt提供了高效便捷的斷點存儲、加載機制。
Snapckpt介紹
Snapckpt是一種為大模型訓練打造的易用、可拓展、高性能的斷點解決方案。Snapckpt提供了高效便捷的斷點存儲、加載機制，其主要有以下幾大優勢：
異步斷點保存：通過優化斷點保存流程，使得斷點保存過程異步化，盡可能減小斷點保存過程對于訓練的中斷，減少訓練階段耗時。
分布式斷點存儲：對模型機優化器參數進行分布式存儲，避免集中式存儲帶來的額外通信開銷，大幅降低存儲與加載斷點的耗時。
多框架支持：支持Megatron-LM、Deepspeed兩大主流大模型訓練框架。
簡單易用：安裝及使用方式簡單，用戶僅需簡單步驟即可使用，提升使用效率。
背景信息
在大規模分布式模型訓練過程中，系統可能因硬件故障或軟件異常導致訓練中斷。為確保訓練進度可恢復，業界普遍采用周期性保存檢查點(Checkpoint)的方案。值得注意的是，檢查點操作耗時與模型參數量呈正相關關系，當面對參數量達百億甚至千億級別的大模型時，每次檢查點保存往往需要耗費數分鐘至十余分鐘不等。特別是在使用Megatron-LM框架或原生PyTorch進行訓練時，檢查點保存過程會強制暫停訓練任務，導致寶貴的計算資源被閑置。因此，開發高效的檢查點機制以降低時間成本和資源浪費，已成為當前大規模模型訓練亟待解決的關鍵問題。
Snapckpt采用多階段異步斷點存儲機制，降低斷點存儲耗時，減少訓練中斷帶來的影響，從而提升訓練速度，提升計算資源有效利用率。
使用方法
1、安裝Snapckpt
Snapckpt Python包下載地址：
下載后進入到工程根目錄
安裝命令：
python3 setup.py build && python3 setup.py install
至此，Snapckpt的軟件包安裝完成，可以開始使用（具體使用方法見“Snapckpt使用方法”章節）。
2、Snapckpt使用方法
Megatron模式使用方法
在使用Megatron-LM訓練框架訓練模型時，您可以使用 snapckpt.megatron_dist_ckpt.save_checkpoin和 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_checkpointDeepspeed模式使用方法
首先，需要初始化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)
性能參考
以Deepseek R1 Distill Qwen 14B模型，單機八卡運行環境為例，相較于同步斷點保存而言，Snapckpt的性能提升測試參考如下：
| 模式 | 全參微調 | LoRA微調 | 預訓練 |
| Megatron | 24倍 | 10倍 | 22倍 |
| Deepspeed | 19倍 | 5倍 | 12倍 |
