使用方式

原生pytorch checkpoint

在训练代码内需要保存checkpoint的py文件内，进行加载和checkpointer初始化：

from eagle.framework_train.pytorch.async_checkpoint.ddp_checkpointer import DdpCheckpointer
checkpointer = DdpCheckpointer("/tmp/flash_checkpoint_example_ckpt")

将pytorch代码里用来存储/读取checkpoint的代码，例如torch.save、torch.load，修改为

checkpointer.save_checkpoint(step, state_dict)
checkpointer.load_checkpoint()

可以参考工程内的示例：examples/ddp_example.py

CTFlashCkpt扩展使用:

英伟达Megatron-LM

按照上述安装方法安装CTFlashCkpt后，修改的英伟达Megatron-LM的默认存储已经被CTFlashCkpt替换，只要按照标准Megatron的pretrain_gpt.py训练，得到的checkpoint即是异步存储的。

华为ModelLink

按照上述安装方法安装CTFlashCkpt后，修改的ModelLink的默认存储已经被CTFlashCkpt替换，只要按照标准ModelLink的pretrain_gpt.py训练，得到的checkpoint即是异步存储的。

镜像使用

平台的公共镜像中我们已经分别为安装了英伟达Megatron-LM和昇腾ModelLink的镜像集成好了CTFlashCkpt的功能，可以直接运行megatrong/modellink：

• 英伟达镜像：ctflashtckpt:v0.1.0-pytorch2.0.0-cuda12.2-ubuntu22.04-x64-llama
• 昇腾镜像：ctflashtckpt:v0.1.0-pytorch2.1.0-cann7.0.0-ubuntu22.04-aarch64-llama

我们提供了运行训练llama的示例，可以自行阅读并修改相关配置进行训练：

• 可参考examples/llama文件夹的run_llama.sh文件，这是训练启动程序；
• var_example.sh，这是启动run_llama.sh的参数示例（不要运行此文件）。

使用效果观察

昇腾 ModelLink的CTFlashCkpt写加速

英伟达 megatron的CTFlashCkpt写加速

注：昇腾及英伟达的测试数据来自于各自独立的测试环境，两组数据相互之间不具备可比性。