已合并
Add test for _EmptyStateDictLoadPlanner on NPU #42763
costan创建于 7月25日
Add test for _EmptyStateDictLoadPlanner on NPU #42763
已合并
Pull Request已成功合入, 合并人@ascend-robot
(感谢 costan 的贡献)7月25日 关联了issue:【社区任务】7月社区任务第三期-Ascend for PyTorch API 一致性开发(38)
atomgit-bot
7月25日 评论:
7月25日 评论:
变更摘要
该 PR 为 torch.distributed.checkpoint.default_planner._EmptyStateDictLoadPlanner API 新增独立的单元测试文件,弥补 PyTorch 官方社区缺少该 API 直接测试的空白。测试覆盖了 keys 参数传参与不传参的正常加载路径,以及 set_up_planner 方法的异常场景(非空 state_dict 和 metadata 为 None 的情况)。该 API 为纯 Python 元数据操作,不涉及 NPU 适配代码修改。
主要改动
- 新增测试文件
test/distributed/checkpoint/test_empty_state_dict_load_planner.py:基于TestCase和分布式环境(gloo 后端、单进程)搭建测试框架,共包含 4 个测试用例。 test_load_keys_none:验证不传keys参数时_EmptyStateDictLoadPlanner能正常加载保存的 checkpoint,并断言"model"键存在于加载后的 state_dict 中。test_load_keys_specified:验证传入keys=["model"]时_EmptyStateDictLoadPlanner能按指定键加载,保存包含model和opt的 state_dict 后仅加载model键。test_set_up_planner_error_non_empty与test_set_up_planner_error_metadata_none:分别验证set_up_planner在传入非空 state_dict 以及 metadata 为None时正确抛出AssertionError,覆盖异常处理路径。


ascend-robot
7月25日 评论:
7月25日 评论:
atomgit-bot
7月25日 评论:
7月25日 评论:
代码审查
审查结束 — 总结
已审查文件:
test/distributed/checkpoint/test_empty_state_dict_load_planner.py:发现 2 个问题
按优先级统计:
- P2:1 个 — 模块级别
dist.init_process_group替代@with_comms装饰器,无清理机制且后端选择不当 - P3:1 个 —
test_load_keys_specified缺少对非指定 key 的排除断言
总体风险判断:中等。该测试文件是新添加的,不影响生产代码。但模块级别的分布式初始化方式存在进程组泄漏和端口冲突风险,在集成测试环境中可能导致不稳定。建议在合入前将分布式初始化改为 @with_comms 装饰器模式,以保持与同目录其他测试文件一致。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 0 |
| 🟡 建议 | 1 |
💬 仅评论


此处折叠了258条消息 查看更多
23 天前 添加了label:approvedlgtm
23 天前 合入了pull request
ascend-robot
23 天前 评论:
23 天前 评论:
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


ascend-robot
23 天前 评论:
23 天前 评论:
流水线 pytorch_gitcode_PR_multiVersion#13995 [ commitID:ece1b151 ] 已完成


【合入来源】
关联任务issue: https://gitcode.com/Ascend/pytorch/issues/3146
【修改方案】
torch.distributed.checkpoint.default_planner._EmptyStateDictLoadPlanner
PyTorch 官方社区在 test/distributed/composable/test_composability/test_pp_composability.py 中有间接测试,但未针对该 API 编写直接独立测试,因此自行编写用例并提交到 test 目录。
该 API 为纯 Python 元数据操作,不涉及 NPU 适配,不需要代码修改。
新增 test/distributed/checkpoint/test_empty_state_dict_load_planner.py,验证 EmptyStateDictLoadPlanner 的加载功能及异常处理
【资料变更】
涉及。上游 PyTorch 在 test/distributed/composable/test_composability/test_pp_composability.py 中有相关测试,但仅作为 PP+DCP 集成测试的工具使用,未针对 EmptyStateDictLoadPlanner 公开 API 编写独立测试。本次提交补充了 15个测试方法,覆盖 keys 参数传参与不传参、异常场景的核心行为,测试文件位于 test/distributed/checkpoint/test_empty_state_dict_load_planner.py。
【接口变更】
不涉及
【功能介绍】
torch.distributed.checkpoint.default_planner._EmptyStateDictLoadPlanner 是 PyTorch 分布式检查点系统的内部工具类,继承自 DefaultLoadPlanner。用于从空 state_dict 加载 DCP 检查点,通过读取元数据重建 state_dict 结构,适用于未初始化模型时加载检查点的场景(如 DCP checkpoint 转 Torch save 文件)。
【功能验证】
test_empty_state_dict_load_planner.py (_EmptyStateDictLoadPlanner)
运行命令:
python -m pytest /workspace/user_data/pytorch-v2.7.1/test/distributed/checkpoint/test_empty_state_dict_load_planner.py -x
运行结果:
【CheckList】