已合并
[test] Add _save_state_dict Api Coverage #42814
lihaokun-2026创建于 26 天前
[test] Add _save_state_dict Api Coverage #42814
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 lihaokun-2026 的贡献）
26 天前 创建了 pull request，commit 8bddc816
26 天前 关联了issue：【社区任务】7月社区任务第三期-Ascend for PyTorch API 一致性开发（59）
atomgit-bot
26 天前 评论：
26 天前 评论：
变更摘要
本 PR 新增了一个 NPU 兼容性测试文件 test_state_dict_saver__save_state_dict.py，用于对 torch.distributed.checkpoint.state_dict_saver._save_state_dict 和 _load_state_dict 在 no_dist 模式下的 NPU 行为进行独立、直接的验证。测试覆盖了基础 NPU tensor 的保存与恢复、多 dtype 支持、标量/空张量、非连续张量、嵌套 state_dict、nn.Module 与 optimizer 状态、CPU 与 NPU 之间的跨设备加载、混合设备 state_dict、自定义 SavePlanner 流程以及 use_collectives=False 场景。
主要改动
- 新增测试类
TestSaveStateDictNpu：继承自TestCase，封装了 NPU 设备获取、同步、_save_state_dict/_load_state_dict调用及 tensor 比较等通用辅助方法。 - 新增辅助类
TinyNpuModule与RecordingSavePlanner：TinyNpuModule是一个小型nn.Module，用于验证 NPU 模型 state_dict 的保存/恢复；RecordingSavePlanner继承自DefaultSavePlanner，记录 planner 各回调调用顺序以验证自定义 planner 流程。 - 覆盖 11 个测试场景：包括基础 NPU tensor 保存加载、
float32/float16/int32/bool/bfloat16多 dtype、标量与空张量、非连续张量（transpose与strided_slice）、嵌套 state_dict、nn.Modulestate_dict、模型+AdamWoptimizer state_dict、NPU→CPU 与 CPU→NPU 跨设备加载、混合 CPU/NPU state_dict、自定义RecordingSavePlanner流程、以及use_collectives=False参数传递。
ascend-robot
26 天前 评论：
26 天前 评论：
atomgit-bot
26 天前 评论：
26 天前 评论：
代码审查
经过详尽的审查，该文件在以下方面均属干净：
- 逻辑正确性：所有张量形状一致，save/load 往返验证完整，
RecordingSavePlanner中的签名自省兼容层设计正确。 - 安全性：无硬编码凭据，无 eval/exec/pickle，无命令注入，无 SQL。
- 资源管理：
tempfile.TemporaryDirectory作为上下文管理器妥善使用；FileSystemWriter/FileSystemReader遵循现有仓库惯例，无需显式调用 close。 - 异常处理：无裸露的
except:块，无被吞掉的异常。 - 可变默认参数：无 — 所有默认参数均为
None、True或False。 - 测试隔离：
torch.npu.set_device(0)遵循整个仓库分布式测试的通行模式。 - 版本门控：
test_save_load_npu_tensor_with_collectives_disabled_no_dist中的提前return以及_call_save_state_dict/_call_load_state_dict中的条件use_collectives参数传递均为有意为之的兼容性保护措施。
关闭摘要
- 按优先级统计发现数量：P0: 0, P1: 0, P2: 0, P3: 0
- 整体风险判断：低。该 diff 新增了一个结构良好、覆盖全面的测试文件，用于在 NPU 上验证
torch.distributed.checkpoint.state_dict_saver._save_state_dict。未发现正确性、安全性、可靠性或破坏性变更问题。
已审查的文件：
test/distributed/checkpoint/test_state_dict_saver__save_state_dict.py— 无问题。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 0 |
| 🟡 建议 | 1 |
💬 仅评论
此处折叠了271条消息 查看更多
chenrayray
7 天前 评论：
7 天前 评论：
/approve
7 天前 添加了label：approvedlgtm
7 天前 删除了label：ci-pipeline-passed
7 天前 合入了pull request
ascend-robot
7 天前 评论：
7 天前 评论：
流水线 pytorch_gitcode_PR_multiVersion#14007 [ commitID：ceb41255 ] 已完成
【合入来源】
#3167
torch.distributed.checkpoint.state_dict_saver._save_state_dict: 【社区任务】7月社区任务第三期-Ascend for PyTorch API 一致性开发（59）【修改方案】
torch.distributed.checkpoint.state_dict_saver._save_state_dict是 DCP 保存流程的核心内部实现函数。外层的 save() 和旧接口 save_state_dict() 最终都会把参数整理好，然后调用 _save_state_dict() 去真正完成 checkpoint 写入，_save_state_dict()的核心功能是：把已经整理好的 state_dict，按照 planner 生成的保存计划，交给 storage writer 写入 checkpoint，并最终生成 metadata。上游用例检查情况：
检索 PyTorch / Torch-NPU 现有测试后，未找到面向
torch.distributed.checkpoint.state_dict_saver._save_state_dict的独立、直接、最小验证用例，因此按任务要求新增 NPU API 验证用例。新增用例文件：
test_state_dict_saver__save_state_dict.py新增用例覆盖场景：
普通 NPU Tensor 保存test_save_npu_tensor_no_dist嵌套状态字典保存test_save_nested_npu_state_dict_no_dist自定义 SavePlannertest_save_npu_tensor_with_custom_planner_no_dist禁用 collectivestest_save_npu_tensor_with_collectives_disabled_no_dist【资料变更】
不涉及。未修改产品资料、接口说明文档或用户可见文档。
【接口变更】
不涉及。未修改跨代码仓接口，也未修改用户可见接口行为。
【功能验证】
测试环境：
测试方法：
source /usr/local/Ascend/ascend-toolkit/set_env.sh TORCH_DEVICE_BACKEND_AUTOLOAD=1 python3 test/distributed/checkpoint/test_state_dict_saver__save_state_dict.py -v测试结果：
【CheckList】