已开启
【自提单】: PyNative 续训加载权重时重复校验权重文件,校验次数随 PP stage 数增长 #2593
lanshaozuishuai创建于  28 天前
lanshaozuishuai成员
28 天前 创建

Checklist

🐞 问题详细描述

PyNative 断点续训加载 MindFormers 格式权重时,对权重文件的存在校验重复执行,大模型放在共享存储上时续训启动明显变慢。

调用链(mindformers/pynative/trainer/trainer.py Trainer._load_checkpoint):

  1. get_checkpoint_path(checkpoint_path) 解析 checkpoint 目录,内部调用 verify_ckpt_valid 做一次完整校验;
  2. 随后 for stage in models 逐 stage 调用 load_checkpoint,其内部再次调用 get_checkpoint_path → verify_ckpt_valid。

因此完整校验共执行 1 + PP stage 数 次。同时 verify_ckpt_valid 遍历 metadata.json 的 storage_data,对每个参数的每个分片逐一 os.path.exists,未按文件去重,同一个 safetensors 文件会被 stat 多次,单次校验的 stat 次数约为"参数数 × 分片数"。

PyNative 推理加载权重(load_checkpoint(checkpoint=get_checkpoint_path(...)))也存在同样的重复解析。

# mindformers/checkpoint/checkpoint.py  load_checkpoint
check_the_param_for_load_ckpt(checkpoint, network)
checkpoint_dir = get_checkpoint_path(checkpoint)   # 调用方已解析过,这里重复执行 verify_ckpt_valid

详细的环境信息描述

与硬件无关,属于代码逻辑问题。任意 PyNative 训练配置设置 checkpoint.load_path 续训即可触发;开启 pipeline parallel 时重复次数随 stage 数增加。

其他辅助信息

修复思路:load_checkpoint 增加 resolve_checkpoint_path 参数(默认 True,静态图行为不变),PyNative 调用方已解析路径时传 False;verify_ckpt_valid 按文件名去重。

版本信息

mindformers master(473e372fc)

likedislike
Llanshaozuishuai成员
18 天前 修改标题为 “【自提单】: PyNative 续训加载权重时重复校验权重文件,校验次数随 PP stage 数增长”,原标题为“[Bug]: PyNative 续训加载权重时重复校验权重文件,校验次数随 PP stage 数增长”