使用方式

原生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的代碼，例如、，修改為

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寫加速

注：昇騰及英偉達的測試數據來自于各自獨立的測試環境，兩組數據相互之間不具備可比性。