昇腾社区首页
中文
注册

tft_register_save_ckpt_handler

接口功能

注册框架侧临终CheckPoint保存函数。

接口格式

mindio_ttp.tft_register_save_ckpt_handler(callback: Callable, ctx = None):

接口参数

参数

是否必选

说明

取值要求

callback

必选

临终CheckPoint保存函数,大模型训练框架提供。

回调函数,不为空,回调函数入参要求请参见下表1

ctx

可选

回调函数上下文。

默认为空。

表1 回调函数参数

参数

是否必选

说明

取值要求

step

必选

临终CheckPoint保存时MindIO TTP选择出的优化器数据对应的step。

正整数。

save_info

必选

不同优化器保存时对应的rank list。

[
{
"type": Enum,参数类型,Attention=0,Moe=1
“ranks”:list,对应参数类型的rank list
},
]

args

必选

tft_set_step_args设置的参数。

默认为空。

ctx

必选

回调函数上下文。

默认为空。

返回值