已开启
【自提单】: PyNative 续训加载权重时重复校验权重文件,校验次数随 PP stage 数增长 #2593
lanshaozuishuai创建于 28 天前
Llanshaozuishuai
18 天前 修改标题为 “【自提单】: PyNative 续训加载权重时重复校验权重文件,校验次数随 PP stage 数增长”,原标题为“[Bug]: PyNative 续训加载权重时重复校验权重文件,校验次数随 PP stage 数增长”
18 天前 修改标题为 “【自提单】: PyNative 续训加载权重时重复校验权重文件,校验次数随 PP stage 数增长”,原标题为“[Bug]: PyNative 续训加载权重时重复校验权重文件,校验次数随 PP stage 数增长”
Checklist
🐞 问题详细描述
PyNative 断点续训加载 MindFormers 格式权重时,对权重文件的存在校验重复执行,大模型放在共享存储上时续训启动明显变慢。
调用链(
mindformers/pynative/trainer/trainer.pyTrainer._load_checkpoint):get_checkpoint_path(checkpoint_path)解析 checkpoint 目录,内部调用verify_ckpt_valid做一次完整校验;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)