已合并
Add test for _EmptyStateDictLoadPlanner on NPU #42763
Add test for _EmptyStateDictLoadPlanner on NPU #42763
已合并
costan创建于 7月25日
costan
costan
7月25日

【合入来源】

关联任务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)

  • test_init_without_keys — 不传 keys 构造 planner, keys 为 None
  • test_init_with_keys — 传入 keys=["model"] ,确认被保存
  • test_set_up_planner_with_non_empty_state_dict — state_dict 非空时 set_up_planner 抛 AssertionError
  • test_set_up_planner_without_metadata — metadata=None 时抛 AssertionError
  • test_set_up_planner_success — 空 state_dict + 合法 metadata,重建出正确 shape/dtype 的张量并保存 metadata
  • test_set_up_planner_with_is_coordinator — 接受并保存 is_coordinator 标志
  • test_keys_none_loads_all — keys=None 时加载 metadata 中全部 key
  • test_keys_filter_loads_subset — keys={"x","z"} 只加载指定子集
  • test_keys_filter_loads_nothing_when_no_match — keys 无匹配时 state_dict 保持为空
  • test_keys_filter_with_planner_data — 通过 planner_data 嵌套路径( model.layer.weight )匹配并构建嵌套 dict
  • test_create_local_plan_after_setup — setup 后 create_local_plan 返回含 TENSOR ReadItem 的 LoadPlan
  • test_create_local_plan_empty_when_keys_filter_all — key 全被过滤时 local plan 为空
  • test_create_global_plan — 单个 rank 时返回 1 个 LoadPlan
  • test_create_global_plan_multiple_ranks — 多 rank 返回对应数量 plan
  • test_finish_plan_passthrough — finish_plan 恒等返回原 plan

运行命令:
python -m pytest /workspace/user_data/pytorch-v2.7.1/test/distributed/checkpoint/test_empty_state_dict_load_planner.py -x
运行结果:

=============================================== test session starts ================================================
platform linux -- Python 3.12.13, pytest-9.1.1, pluggy-1.6.0
rootdir: /workspace/user_data/pytorch-v2.7.1
configfile: pyproject.toml
plugins: hypothesis-6.156.6
collected 15 items                                                                                                 

user_data/pytorch-v2.7.1/test/distributed/checkpoint/test_empty_state_dict_load_planner.py ...............   [100%]

================================================ 15 passed in 4.03s ================================================

【CheckList】

PR提交人对以下CheckList自检项进行全量自检,自检通过或不涉及,均修改 [ ] 为 [x]

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 costan 的贡献)
costancostan
7月25日 创建了 pull request,commit c10f8f3e
costancostan
7月25日 关联了issue:【社区任务】7月社区任务第三期-Ascend for PyTorch API 一致性开发(38)
atomgit-bot
atomgit-bot
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 能按指定键加载,保存包含 modelopt 的 state_dict 后仅加载 model 键。
  • test_set_up_planner_error_non_emptytest_set_up_planner_error_metadata_none:分别验证 set_up_planner 在传入非空 state_dict 以及 metadata 为 None 时正确抛出 AssertionError,覆盖异常处理路径。
likedislike
atomgit-bot
atomgit-bot
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

💬 仅评论

likedislike
此处折叠了258条消息 查看更多
renyujin成员
23 天前 评论:

/lgtm
/approve

likedislike
ascend-robotascend-robot成员
23 天前 添加了label:approvedlgtm
ascend-robotascend-robot成员
23 天前 合入了pull request
ascend-robot
ascend-robot成员
23 天前 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

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