v2.9.0

文件	最后提交记录	最后更新时间
test_StorageWriter_api.py	test(distributed/checkpoint): Add StorageWriter storage_meta API tests Co-authored-by: Jwerr<19166910919@163.com> # message auto-generated for no-merge-commit merge: !35307 merge test-StorageWriter-api-v2.9.0 into v2.9.0 test(distributed/checkpoint): Add StorageWriter storage_meta API tests Created-by: Jwerr Commit-by: Jwerr Merged-by: ascend-robot Description: <!-- PR描述模板更新日期：20260203 --> # 【合入来源】 关联issue #1895 PyTorch 社区缺少对 torch.distributed.checkpoint.StorageWriter.storage_meta（实际类路径：torch.distributed.checkpoint.storage.StorageWriter.storage_meta）的直接 API 级测试，且未覆盖 NPU。本 PR 在 test/distributed/checkpoint/test_StorageWriter_api.py 补齐单卡 + 多卡 NPU 用例。 - [ ] 需求 - [x] 问题单 - [ ] issue/工单 - [ ] 重构优化 - [ ] 资料更新 # 【修改方案】 一、API功能说明 torch.distributed.checkpoint.StorageWriter.storage_meta是 torch.distributed.checkpoint.StorageWriter 基类下的可选钩子，父类默认返回 None，子类可按需覆写并返回 torch.distributed.checkpoint.metadata.StorageMeta。 torch.distributed.checkpoint.metadata.StorageMeta 用于描述 checkpoint 存储元信息，包含 checkpoint_id、save_id、load_id、modules 等字段。其中 checkpoint_id 表示存储位置 ID，save_id 表示本次保存 UUID，modules 默认为空列表。 在保存流程中，writer.storage_meta() 返回的 torch.distributed.checkpoint.metadata.StorageMeta 会通过 torch.distributed.checkpoint.save 透传给 torch.distributed.checkpoint.planner.SavePlanner.set_up_planner(state_dict, storage_meta=...)，用于保存规划阶段感知 writer 侧元数据。 二、测试文件 test_StorageWriter_api.py 验证内容 1. 基类默认行为验证 定义 _MinimalStorageWriter 作为 torch.distributed.checkpoint.storage.StorageWriter 最小实现类，仅实现抽象方法，不覆写 torch.distributed.checkpoint.storage.StorageWriter.storage_meta，验证： * torch.distributed.checkpoint.storage.StorageWriter.storage_meta 可调用； * 父类默认 storage_meta() 返回 None。 2. FileSystemWriter 覆写行为验证 通过 torch.distributed.checkpoint.FileSystemWriter 验证真实 writer 的 storage_meta 返回值： * 返回值类型为 torch.distributed.checkpoint.metadata.StorageMeta； * checkpoint_id 与初始化 checkpoint 目录一致； * save_id 为合法 UUID； * modules 默认为空列表； * reset(checkpoint_id=new_dir) 后 checkpoint_id 更新，save_id 变化。 3. 保存流程验证 通过 torch.distributed.checkpoint.save 验证 storage_meta 在实际保存流程中的可用性： * 使用 NPU Tensor 执行 torch.distributed.checkpoint.save(..., no_dist=True) 后，writer.storage_meta() 仍返回合法 torch.distributed.checkpoint.metadata.StorageMeta； * 空 state_dict 保存后，writer.storage_meta() 仍保持有效； * save -> reset -> save 后 save_id 发生变化，验证 writer 生命周期更新有效。 4. SavePlanner 透传验证 定义 _CapturingPlanner 继承 torch.distributed.checkpoint.default_planner.DefaultSavePlanner，在 set_up_planner 中捕获 storage_meta，验证： * torch.distributed.checkpoint.save 会将 writer.storage_meta() 返回值透传给 torch.distributed.checkpoint.planner.SavePlanner.set_up_planner； * planner 捕获到的 checkpoint_id、save_id 与 writer.storage_meta() 保持一致。 5. 多 NPU 分布式验证 TestStorageMetaDistributed 继承 torch.testing._internal.distributed._shard.sharded_tensor.ShardedTensorTestBase，通过 torch_npu.testing.common_distributed.with_comms 初始化 HCCL 进程组，并通过 torch_npu.testing.common_distributed.skipIfUnsupportMultiNPU(2) 限定 2 卡 NPU 环境运行，验证： * rank 0 创建共享 checkpoint 目录，并通过 torch.distributed.broadcast_object_list 广播给其他 rank； * 分布式 torch.distributed.checkpoint.save 后，各 rank 的 checkpoint_id 通过 torch.distributed.all_gather_object 收集并保持一致； * 分布式 save -> reset -> save 后 save_id 发生变化； * 分布式保存后 writer.storage_meta() 仍返回 torch.distributed.checkpoint.metadata.StorageMeta，且 save_id 为合法 UUID。 三、NPU适配 torch.distributed.checkpoint.storage.StorageWriter.storage_meta 和 torch.distributed.checkpoint.metadata.StorageMeta 均属于 Python 层 checkpoint 元数据抽象，不涉及 NPU 算子、NPU kernel、设备内存管理或 HCCL 通信协议本身，因此 API 本身无需针对 NPU 修改。 本 PR 主要在测试用例层面验证 NPU 场景可用性： * 单卡用例继承 torch_npu.testing.testcase.TestCase，使用 torch.Tensor.npu() 构造 NPU Tensor 后进入 torch.distributed.checkpoint.save 流程； * 多卡用例通过 torch_npu.testing.common_distributed.with_comms 初始化 HCCL 进程组； * 通过 torch_npu.testing.common_distributed.skipIfUnsupportMultiNPU(2) 限定多 NPU 环境，避免设备数量不足时误失败； * 通过 torch.distributed.broadcast_object_list 保证各 rank 使用同一 checkpoint 目录； * 通过 torch.distributed.all_gather_object 验证 checkpoint_id 跨 rank 一致； * 测试结束后通过 destroy_pg 释放分布式进程组资源。 # 【资料变更】 不涉及 # 【接口变更】 不涉及 # 【功能验证】 在2.7.1 2.9.0 2.10.0 2.11.0 2.12.0以及master版本上执行该用例，均通过，本分支对应2.9.0日志（完整版在评论区）如下： `` /root/work/test_StorageWriter_api.py:303: FutureWarning: save_state_dict is deprecated and will be removed in future versions.Please use save instead. save_state_dict( . ---------------------------------------------------------------------- Ran 8 tests in 178.952s OK (skipped=6) `` # 【CheckList】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [x] 代码注释完备，正确记录错误日志 - [x] 代码实现进行了返回值、空指针等校验 - [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [x] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!35307	3 个月前
test_broadcasting_reader.py	test(distributed):add test for validate_checkpoint_id, reset, set_up_storage_reader from BroadcastingTorchSaveReader for v2.9.0 Co-authored-by: xh-zhan<zhanxuhan@sjtu.edu.cn> # message auto-generated for no-merge-commit merge: !33950 merge test-v2.9.0 into v2.9.0 test(distributed):add test for validate_checkpoint_id, reset, set_up_storage_reader from BroadcastingTorchSaveReader for v2.9.0 Created-by: xh-zhan Commit-by: xh-zhan Merged-by: ascend-robot Description: /kind feature ### What does this PR do / why do we need it: 本次 PR 为 BroadcastingTorchSaveReader 类中的三个核心 API 新增了单元测试用例，以确保其在 torch_npu 环境下的行为与 PyTorch 原生行为保持一致。 测试用例路径：test/distributed/checkpoint/test_broadcasting_reader.py ### API 功能说明: - **torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.validate_checkpoint_id**：纯逻辑 API，用于检查给定的 checkpoint 文件路径是否存在。 - **torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.reset**：纯逻辑 API，用于重置 BroadcastingTorchSaveReader 实例的 checkpoint_id 属性，支持更新为新的路径或 None。 - **torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.set_up_storage_reader**：用于设置存储读取器的元数据及协调者（coordinator）状态。包含协调者身份校验和 checkpoint_id 非空检查。 ### 用例完备性说明: 由于Pytorch社区无上述三个API的测试用例，故本次PR新增用例test_broadcasting_reader.py。本次提交的测试用例覆盖了以下场景，确保 API 的健壮性： - validate_checkpoint_id: - ✅ 存在的文件路径 → 返回 True - ✅ 不存在的文件路径 → 返回 False - reset: - ✅ 更新为新路径 → checkpoint_id 正确变更 - ✅ 更新为 None → checkpoint_id 正确变更为 None - set_up_storage_reader: - ✅ 正确设置 is_coordinator 属性 - ✅ 协调者 rank 不匹配时抛出 AssertionError - ✅ checkpoint_id 为 None 时抛出 AssertionError > 注：测试中使用 unittest.mock 模拟 torch.distributed.get_rank，以在单进程环境下验证分布式协调者校验逻辑，避免对真实多进程环境的依赖。 ### 适配方案: 经分析，这三个 API 均为纯 Python 逻辑，不依赖任何 NPU 特定算子或硬件特性。本次提交 **无需对 API 本身进行任何代码适配**，仅增加单元测试以完成 NPU 环境下的兼容性验证。 ### Which issue(s) this PR fixes: Fixes [#1680](https://gitcode.com/Ascend/pytorch/issues/1680) ### Special notes for your reviewers: 1. **测试文件路径**：test/distributed/checkpoint/test_broadcasting_reader.py 2. **本地验证结果**：所有测试用例已在昇腾 NPU 环境下成功运行，且本地2.6.0，2.7.1，2.8.0，2.9.0，2.10.0，2.11.0版本下都运行成功，输出如下： Ran 7 tests in 0.011s OK 3. **命名规范**：测试文件命名符合 torch-npu 仓库规范（test_ 开头 + 核心关键词），文件内添加了完整 API 名称和测试维度注释，便于后续检索和维护。 4. **需合入的分支**：本PR合入v2.9.0 分支。 See merge request: Ascend/pytorch!33950	3 个月前
test_broadcasting_reader_api.py	test: add NPU tests for BroadcastingTorchSaveReader APIs Co-authored-by: zj<2943168611@qq.com> # message auto-generated for no-merge-commit merge: !35057 merge test/2.9.0-broadcasting-reader-api-npu into v2.9.0 test: add NPU tests for BroadcastingTorchSaveReader APIs Created-by: zjucn Commit-by: zj Merged-by: ascend-robot Description: <!-- PR描述模板更新日期：20260203 --> # 【合入来源】 - [x] 问题单：https://gitcode.com/Ascend/pytorch/issues/1883 本 PR 覆盖以下 API： - torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader - torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_metadata - torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_local_plan - torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_global_plan - torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_data # 【修改方案】 ## 一、解决方案 针对 issue 中提出的 BroadcastingTorchSaveReader API 专项测试缺失问题，本 PR 新增 test_broadcasting_reader_api.py，以“API 级行为验证 + NPU 分布式数据面验证”的方式补齐相关测试。 本 PR 做如下修改： - 新增 test_broadcasting_reader_api.py，覆盖 BroadcastingTorchSaveReader 的 metadata、plan 透传、数据加载和错误边界行为。 - 使用 with_comms 和 skipIfUnsupportMultiNPU(2) 验证 read_data 在 NPU/HCCL 分布式环境下的 broadcast 加载路径。 - 使用 DynamicMetaLoadPlanner 覆盖普通 torch save 文件无 DCP metadata 时的推荐加载组合。 - 不修改 torch-npu 生产代码，不改变 API 签名、返回结构或运行时语义。 新增测试文件： text test/distributed/checkpoint/test_broadcasting_reader_api.py ## 二、用例完备性说明 本次新增 6 个测试用例，覆盖 BroadcastingTorchSaveReader 的生命周期入口、计划透传、数据加载主路径和错误边界。 ### 1. test_read_metadata_returns_empty_metadata 覆盖 API：read_metadata 验证内容： - 返回对象是 Metadata。 - state_dict_metadata 为空字典。 - 不需要真实 checkpoint 文件，也不依赖分布式环境。 该用例证明 BroadcastingTorchSaveReader 在 torch save 无 metadata 文件的情况下，能按设计返回空 metadata，并把 metadata 构造职责交给 planner。 ### 2. test_prepare_local_plan_returns_input_plan 覆盖 API：prepare_local_plan 验证内容： - 构造带 storage_data 和 planner_data 的 LoadPlan。 - 调用 prepare_local_plan 后返回值与输入对象是同一个对象。 - 输入计划中的附加数据不会被替换或丢失。 该用例证明 reader 不会修改本地 planner 结果，符合该 API 的 no-op 设计。 ### 3. test_prepare_global_plan_returns_input_plans 覆盖 API：prepare_global_plan 验证内容： - 构造多个 rank 的 LoadPlan 列表。 - 调用 prepare_global_plan 后返回值与输入列表是同一个列表。 - 各 rank 的 plan 不被重排、不被替换。 该用例证明 reader 在全局 planning 阶段不引入额外调度逻辑，保持 DCP 默认计划流转。 ### 4. test_read_data_rejects_byte_io_items 覆盖 API：read_data 验证内容： - 手工构造 LoadItemType.BYTE_IO 类型的 ReadItem。 - 使用真实 torch save 文件路径初始化 reader。 - 调用 read_data 时断言抛出 only supports loading Tensors。 该用例覆盖官方端到端成功路径没有覆盖的错误边界，证明 reader 对非 tensor 对象的限制是显式、可预期的。 ### 5. test_read_data_loads_torch_save_tensor_to_npu_state_dict 覆盖 API：read_data 验证内容： - rank 0 使用 torch.save 保存普通 tensor checkpoint。 - 使用 with_comms 启动 NPU/HCCL 分布式环境。 - 每个 rank 构造 NPU 目标 tensor。 - 通过 dcp.load + DynamicMetaLoadPlanner + BroadcastingTorchSaveReader 加载。 - 校验加载后的 NPU tensor 与源 tensor 一致。 该用例覆盖 read_data 的核心成功路径，证明 torch save 文件可以经 coordinator 读取、broadcast 分发并写回 NPU state_dict。 ### 6. test_read_data_handles_nested_state_dict 覆盖 API：read_data 验证内容： - 使用嵌套结构 {"model": {"weight": ..., "bias": ...}} 保存 torch save checkpoint。 - 目标 state_dict 使用同样嵌套结构，但 tensor 分配在 NPU。 - 使用 DynamicMetaLoadPlanner 默认 flatten 行为加载。 - 校验嵌套 key 下的 weight 和 bias 均正确写回。 该用例覆盖 read_data 中 flatten_state_dict 为 True 时的关键路径，证明嵌套 state_dict 在 NPU 分布式加载场景下可用。 ### 覆盖关系总结 - metadata 阶段：test_read_metadata_returns_empty_metadata - local plan 阶段：test_prepare_local_plan_returns_input_plan - global plan 阶段：test_prepare_global_plan_returns_input_plans - read_data 错误边界：test_read_data_rejects_byte_io_items - read_data 普通 tensor 成功路径：test_read_data_loads_torch_save_tensor_to_npu_state_dict - read_data 嵌套/flatten 成功路径：test_read_data_handles_nested_state_dict 上述用例覆盖了 reader 生命周期中本 PR 涉及的全部 API，并将纯 Python 计划行为与真实 NPU 通信加载行为分开验证，便于定位问题。 ## 三、NPU 适配说明 本次不修改 torch-npu 生产代码，仅新增测试用例。 适配验证策略： - 使用真实 NPU tensor 作为目标 state_dict，验证 reader 写回路径能落到 NPU 设备。 - 使用 with_comms 和 skipIfUnsupportMultiNPU(2) 启动 HCCL 分布式环境，覆盖 dist.broadcast 的 NPU 通信路径。 - torch save 文件由 rank 0 写入，其他 rank 通过 barrier 等待，贴近 BroadcastingTorchSaveReader 的 coordinator 读取模型。 - 使用 DynamicMetaLoadPlanner 复现官方推荐组合，覆盖无 DCP metadata 的 torch save 加载场景。 经检查，2.7.1、2.9.0、2.10.0、2.11.0、2.12.0 和 master 中本测试依赖的 BroadcastingTorchSaveReader 核心语义、LoadPlan / ReadItem 字段以及 dcp.load 调用方式保持兼容。高版本主要是类型标注和内部异常写法调整，不影响本用例验证的 metadata、plan 透传、BYTE_IO 错误边界和 tensor broadcast 加载路径，因此不同版本不需要拆分修改测试逻辑。 为什么不需要修改 API： - read_metadata、prepare_local_plan、prepare_global_plan 是 DCP storage reader 生命周期的通用 Python 接口，不依赖硬件后端。 - read_data 中与硬件相关的部分是 tensor device 和 dist.broadcast，这部分通过真实 NPU tensor 和 HCCL 进程组验证即可。 - 该 reader 的行为约束来自 PyTorch DCP 通用设计，NPU 适配目标是验证现有语义在 NPU 下可用，而不是改变 API 签名或返回结构。 # 【资料变更】 > 不涉及 # 【接口变更】 > 不涉及 # 【功能验证】 在 2.7.1、2.9.0、2.10.0、2.11.0、2.12.0, master版本上执行新增测试，均通过。 ---------------------------------------------------------------------- Ran 6 tests in 1.367s OK # 【CheckList】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [x] 代码注释完备，正确记录错误日志 - [x] 代码实现进行了返回值、空指针等校验 - [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [x] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!35057	3 个月前
test_checkpoint.py	【fix】修复assertError类型报错的问题 Co-authored-by: xiaoqi-zhou<zhouxiaoqi9@h-partners.com> # message auto-generated for no-merge-commit merge: !31687 merge assert_v2.9.0 into v2.9.0 【fix】修复assertError类型报错的问题 Created-by: xiaoqi-zhou Commit-by: xiaoqi-zhou Merged-by: ascend-robot Description: <!-- PR描述模板更新日期：20260203 --> # 【合入来源】 问题单 - [ ] 需求 - [x] 问题单 - [ ] issue/工单 - [ ] 重构优化 - [ ] 资料更新 # 【修改方案】 在torch2.9.0及以后的版本中，运行test/distributed/checkpoint/test_checkpoint.py中的_test_dist_failure方法，错误场景从ValueError 扩展为ValueError 和 AssertionError。 使用assert判断时需要将错误场景扩大，与开源代码对齐 # 【资料变更】 不涉及 # 【接口变更】 不涉及 # 【功能验证】 本地用例通过： python test_checkpoint.py -v -k test_load_error_handling  # 【CheckList】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [ ] 代码注释完备，正确记录错误日志 - [ ] 代码实现进行了返回值、空指针等校验 - [ ] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [ ] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!31687	5 个月前
test_fsdp_model_state.py	[fix] with_comms can't set self.device in 2.8.0+ bugfix Co-authored-by: zqwenn<zhangqiongwen@huawei.com> # message auto-generated for no-merge-commit merge: !26344 merge v2.9.0_with_comms_fix into v2.9.0 [fix] with_comms can't set self.device in 2.8.0+ bugfix Created-by: zqwenn Commit-by: zqwenn Merged-by: ascend-robot Description: <!-- Thanks for sending a pull request! --> **What type of PR is this?** > Uncomment only one /kind <> line, hit enter to put that in a new line, and remove leading whitespaces from that line: > > /kind bug > /kind task > /kind feature **What does this PR do / why do we need it**: **Which issue(s) this PR fixes**: <!-- *Automatically closes linked issue when PR is merged. Usage: Fixes #<issue number>, or Fixes (paste link of issue). --> Fixes # **Special notes for your reviewers**: See merge request: Ascend/pytorch!26344	9 个月前
test_fsdp_optim_state.py	[fix] with_comms can't set self.device in 2.8.0+ bugfix Co-authored-by: zqwenn<zhangqiongwen@huawei.com> # message auto-generated for no-merge-commit merge: !26344 merge v2.9.0_with_comms_fix into v2.9.0 [fix] with_comms can't set self.device in 2.8.0+ bugfix Created-by: zqwenn Commit-by: zqwenn Merged-by: ascend-robot Description: <!-- Thanks for sending a pull request! --> **What type of PR is this?** > Uncomment only one /kind <> line, hit enter to put that in a new line, and remove leading whitespaces from that line: > > /kind bug > /kind task > /kind feature **What does this PR do / why do we need it**: **Which issue(s) this PR fixes**: <!-- *Automatically closes linked issue when PR is merged. Usage: Fixes #<issue number>, or Fixes (paste link of issue). --> Fixes # **Special notes for your reviewers**: See merge request: Ascend/pytorch!26344	9 个月前
test_loadplan_api.py	test: add load planner api npu tests Co-authored-by: zj<2943168611@qq.com> # message auto-generated for no-merge-commit merge: !35095 merge test/2.9.0-loadplan-api-npu into v2.9.0 test: add load planner api npu tests Created-by: zjucn Commit-by: zj Merged-by: ascend-robot Description: <!-- PR描述模板更新日期：20260203 --> # 【合入来源】 - [x] 问题单：https://gitcode.com/Ascend/pytorch/issues/1885 本 PR 覆盖以下 API： - torch.distributed.checkpoint.LoadPlan - torch.distributed.checkpoint.LoadPlanner - torch.distributed.checkpoint.LoadPlanner.set_up_planner - torch.distributed.checkpoint.LoadPlanner.create_local_plan - torch.distributed.checkpoint.LoadPlanner.create_global_plan - torch.distributed.checkpoint.LoadPlanner.finish_plan - torch.distributed.checkpoint.LoadPlanner.load_bytes - torch.distributed.checkpoint.LoadPlanner.resolve_tensor - torch.distributed.checkpoint.LoadPlanner.commit_tensor # 【修改方案】 ## 一、解决方案 针对 issue 中提出的 LoadPlan / LoadPlanner API 专项测试缺失问题，本 PR 新增 test_loadplan_api.py，以“API 级计划行为验证 + NPU state_dict 加载路径验证”的方式补齐相关测试。 本 PR 做如下修改： - 新增 test_loadplan_api.py，覆盖 LoadPlan / LoadPlanner 的计划生成、计划传递、bytes 加载、tensor 定位和 tensor 提交行为。 - 使用真实 NPU tensor 构造目标 state_dict，验证 planner 在 NPU tensor、NPU tensor view、非连续 NPU tensor 场景下的行为。 - 使用 FileSystemWriter / FileSystemReader 的 no_dist=True 路径，验证真实 checkpoint 读写流程可以加载 tensor 和 bytes 到 NPU state_dict。 - 通过自定义 PlanDataLoadPlanner 验证 planner_data / storage_data 在 local plan、global plan、finish plan 间的传递。 - 通过自定义 MaterializeOnCpuLoadPlanner 验证 resolve_tensor 返回 CPU 临时 tensor 后，commit_tensor 可将数据回写到 NPU 目标 tensor。 - 不修改 torch-npu 生产代码，不改变 API 签名、返回结构或运行时语义。 新增测试文件： text test/distributed/checkpoint/test_loadplan_api.py ## 二、用例完备性说明 本次新增 12 个测试用例，覆盖 LoadPlan / LoadPlanner 加载生命周期、错误边界、NPU 写入路径和跨版本兼容点。 ### 1. test_default_load_planner_local_global_finish_plan 覆盖 API： - LoadPlan - LoadPlanner.set_up_planner - LoadPlanner.create_local_plan - LoadPlanner.create_global_plan - LoadPlanner.finish_plan 验证内容： - NPU tensor 生成 LoadItemType.TENSOR 类型 ReadItem。 - bytes 对象生成 LoadItemType.BYTE_IO 类型 ReadItem。 - tensor ReadItem 的 dest_offsets、storage_offsets 和 lengths 正确。 - DefaultLoadPlanner.create_global_plan 默认透传 local plan。 - DefaultLoadPlanner.finish_plan 默认透传 global plan。 该用例证明默认 planner 可以基于 NPU 目标 state_dict 生成正确 LoadPlan，并保持默认 global / finish plan 语义。 ### 2. test_default_load_planner_creates_multiple_tensor_read_items 覆盖 API： - LoadPlan - LoadPlanner.create_local_plan 验证内容： - 构造包含两个 tensor chunk 的 checkpoint metadata。 - 单个 NPU 目标 tensor 生成两个 tensor ReadItem。 - 每个 ReadItem 的 dest_index、storage_index、storage_offsets 和 lengths 正确。 该用例覆盖 PyTorch 2.11.0 及 master 中 create_read_items_for_chunk_list 算法改写后的公共行为一致性。 ### 3. test_default_load_planner_strict_and_partial_load 覆盖 API： - LoadPlanner.set_up_planner - LoadPlanner.create_local_plan 验证内容： - allow_partial_load=False 时，目标 state_dict 中存在 checkpoint 缺失 key 会抛出 Missing key in checkpoint。 - allow_partial_load=True 时，planner 只为 checkpoint 中存在的 key 生成读取计划。 该用例证明 set_up_planner 记录的目标 state_dict 会参与 strict / partial load 检查。 ### 4. test_default_load_planner_size_mismatch 覆盖 API： - LoadPlanner.set_up_planner - LoadPlanner.create_local_plan 验证内容： - 目标 NPU tensor shape 与 checkpoint metadata shape 不一致时抛出 Size mismatch。 该用例证明目标 NPU tensor 的 shape 会参与 metadata shape 校验。 ### 5. test_resolve_tensor_returns_npu_narrow_view 覆盖 API： - LoadPlanner.resolve_tensor - LoadPlanner.commit_tensor 验证内容： - resolve_tensor 根据 dest_offsets 和 lengths 返回 NPU narrow view。 - 向该 view 写入数据后，原始 NPU tensor 对应切片被正确更新。 - 默认空 commit_tensor 不影响 view 写回结果。 该用例证明默认 planner 可以为 NPU tensor 返回可写入 view。 ### 6. test_resolve_tensor_handles_non_contiguous_npu_target 覆盖 API： - LoadPlanner.resolve_tensor - LoadPlanner.commit_tensor 验证内容： - 目标 tensor 是非连续 NPU tensor。 - resolve_tensor 返回的 view 仍位于 NPU。 - 写入 view 后，非连续目标 tensor 对应区域正确更新。 该用例证明默认 planner 对非连续 NPU tensor 目标仍能正确定位写入位置。 ### 7. test_load_bytes_updates_flattened_original_state_dict 覆盖 API： - LoadPlanner.load_bytes 验证内容： - DefaultLoadPlanner 默认开启 flatten_state_dict。 - bytes ReadItem 使用 flatten FQN。 - load_bytes 反序列化结果后写回原始嵌套 state_dict。 该用例证明 flatten 场景下 bytes 对象可以正确写回原始嵌套结构。 ### 8. test_load_bytes_updates_unflattened_state_dict 覆盖 API： - LoadPlanner.load_bytes 验证内容： - DefaultLoadPlanner(flatten_state_dict=False, flatten_sharded_tensors=False) 关闭 flatten 行为。 - bytes 对象直接写回顶层 state_dict。 该用例证明 unflatten 场景下 bytes 对象可以直接写回目标 key。 ### 9. test_load_state_dict_accepts_custom_plan_data 覆盖 API： - LoadPlan.storage_data - LoadPlan.planner_data - LoadPlanner.create_local_plan - LoadPlanner.create_global_plan - LoadPlanner.finish_plan 验证内容： - 自定义 planner 在 local plan 中写入 planner_data。 - 自定义 planner 在 global plan 中写入 storage_data 并改写 planner_data。 - finish_plan 能接收到 global plan 写入的数据。 - 真实 load_state_dict 流程能完成 NPU tensor 加载。 该用例证明自定义 plan 数据可以经过真实 DCP load 生命周期传递。 ### 10. test_filesystem_metadata_version_when_supported 覆盖 API： - FileSystemReader.read_metadata - Metadata.version 兼容路径 验证内容： - 读取真实 filesystem checkpoint metadata。 - 当当前 PyTorch 版本支持 metadata.version 字段时，校验其等于官方 CURRENT_DCP_VERSION。 - 当当前 PyTorch 版本不支持该字段时，仅验证 metadata 可正常读取。 该用例覆盖 PyTorch 2.9.0 及之后版本新增 metadata version 字段的兼容行为。 ### 11. test_custom_commit_tensor_materializes_cpu_tensor_to_npu 覆盖 API： - LoadPlanner.resolve_tensor - LoadPlanner.commit_tensor 验证内容： - 自定义 planner 的 resolve_tensor 返回 CPU 临时 tensor。 - StorageReader 将 checkpoint 数据写入 CPU 临时 tensor。 - 自定义 commit_tensor 将 CPU 临时 tensor copy 回 NPU 目标 tensor。 - 加载后的 NPU tensor 与保存前数据一致。 该用例证明自定义 commit_tensor 可以支持临时 tensor materialize 后再写回 NPU 的扩展路径。 ### 12. test_filesystem_load_tensor_and_bytes_to_npu_state_dict 覆盖 API： - LoadPlanner.load_bytes - LoadPlanner.resolve_tensor - LoadPlanner.commit_tensor - FileSystemReader - FileSystemWriter 验证内容： - 使用真实 FileSystemWriter 保存包含 NPU tensor 和 bytes 对象的 state_dict。 - 使用真实 FileSystemReader 和 DefaultLoadPlanner 加载到 NPU 目标 state_dict。 - 校验 tensor 数据和 bytes 对象均正确写回。 该用例证明默认 filesystem checkpoint 读写路径可以加载 tensor 和 bytes 到 NPU state_dict。 ### 覆盖关系总结 - plan 生成与透传：test_default_load_planner_local_global_finish_plan - 多 chunk 读计划：test_default_load_planner_creates_multiple_tensor_read_items - strict / partial load：test_default_load_planner_strict_and_partial_load - shape mismatch：test_default_load_planner_size_mismatch - NPU tensor view 写入：test_resolve_tensor_returns_npu_narrow_view - 非连续 NPU tensor 写入：test_resolve_tensor_handles_non_contiguous_npu_target - bytes flatten 写回：test_load_bytes_updates_flattened_original_state_dict - bytes unflatten 写回：test_load_bytes_updates_unflattened_state_dict - 自定义 plan data 传递：test_load_state_dict_accepts_custom_plan_data - metadata version 兼容：test_filesystem_metadata_version_when_supported - 自定义 commit 回写：test_custom_commit_tensor_materializes_cpu_tensor_to_npu - 真实 filesystem load：test_filesystem_load_tensor_and_bytes_to_npu_state_dict 上述用例覆盖了 issue 中列出的全部 API，并将纯 Python 计划行为、NPU tensor 定位写入行为和真实 filesystem 加载路径分开验证，便于定位问题。 ## 三、NPU 适配说明 本次不修改 torch-npu 生产代码，仅新增测试用例。 适配验证策略： - 使用 torch.zeros().to(device_type)、torch.arange().to(device_type) 等方式构造真实 NPU state_dict。 - 使用 FileSystemWriter / FileSystemReader 的 no_dist=True 单进程路径，避免引入多进程 / HCCL 依赖，保持 API 测试稳定。 - 对 DefaultLoadPlanner.resolve_tensor 返回的 NPU view 进行直接写入验证。 - 对非连续 NPU tensor 目标进行直接写入验证。 - 对多 chunk checkpoint metadata 生成的多个 ReadItem 进行直接计划校验，覆盖 resharding 读计划行为。 - 对 2.9.0 及之后版本的 checkpoint metadata version 字段进行兼容验证。 - 通过自定义 MaterializeOnCpuLoadPlanner 覆盖 CPU 临时 tensor 经 commit_tensor 回写 NPU 的扩展路径。 - 使用 NPU tensor 和 bytes 混合 state_dict 验证 FileSystemReader / FileSystemWriter 与默认 planner 的组合加载路径。 经检查，PyTorch 2.7.1、2.9.0、2.10.0、2.11.0 和 master 中 LoadPlan / LoadPlanner 加载侧公共 API 的方法签名保持稳定。版本间主要差异如下： - PyTorch 2.9.0 起 torch.distributed.checkpoint.Metadata 新增 version 字段，FileSystemWriter 会写入 CURRENT_DCP_VERSION。 - PyTorch 2.10.0 将部分内部 assert 改为显式 AssertionError，不影响正常加载路径。 - PyTorch 2.11.0 起 planner_helpers.create_read_items_for_chunk_list 的多 chunk 匹配算法改为 sweep-line 实现，公共行为应保持一致。 - master 主要新增 save plan 校验与缓存相关变化，不改变本 PR 覆盖的加载侧 API 签名。 为什么不需要修改 API： - LoadPlan / LoadPlanner 是 DCP load 阶段的通用计划与数据定位抽象，不绑定具体设备后端。 - 本 PR 验证这些通用接口在目标 state_dict 位于 NPU 时仍能正确生成计划、定位 tensor/view、加载 bytes，并通过自定义 planner 完成 commit 回写。 - NPU 适配目标是验证现有语义在 NPU 下可用，而不是改变 API 签名或返回结构。 # 【资料变更】 > 不涉及 # 【接口变更】 > 不涉及 # 【功能验证】 在 2.7.1、2.9.0、2.10.0、2.11.0、2.12.0以及 master 版本上执行新增测试，均通过。 执行命令： bash python test_loadplan_api.py 结果示例： text ---------------------------------------------------------------------- Ran 12 tests in 1.367s OK # 【CheckList】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [x] 代码注释完备，正确记录错误日志 - [x] 代码实现进行了返回值、空指针等校验 - [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [x] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!35095	3 个月前
test_planner_api.py	test(distributed): add test for WriteItem.tensor_storage_size Co-authored-by: qq_65095566<1922937030@qq.com> # message auto-generated for no-merge-commit merge: !35505 merge test-writeitem-tensor-storage-size-v2.9.0 into v2.9.0 test(distributed): add test for WriteItem.tensor_storage_size Created-by: qq_65095566 Commit-by: qq_65095566 Merged-by: ascend-robot Description: <!-- PR描述模板更新日期：20260203 --> # 【合入来源】 > <font color="red">**[[Usage]: API一致性说明：torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 缺少专项验证用例 #1909](https://gitcode.com/Ascend/pytorch/issues/1909)**</font>\ > <font color="red">**请勿携带内部流程信息（需求链接、问题单、内部issue等）**</font> - [ ] 需求 - [x] 问题单 - [ ] issue/工单 - [ ] 重构优化 - [ ] 资料更新 # 【修改方案】 > 请描述修改内容的具体实现，涉及哪些组件之间进行交互，可以用1、2、3、...进行罗列\ > 如果是需求或者重构类的PR，需要补充详细设计文档（说明上下游组件关系、时序图、类图、DFX能力等内容） 本 PR 针对 torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 新增专项验证用例。 一、API 功能说明 torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 是 distributed checkpoint planner 中 WriteItem 提供的元数据辅助方法，用于返回当前写入项对应 tensor 的存储大小，单位为 bytes。 该 API 的核心行为如下： 1. 当 WriteItem.tensor_data 为 None 时，表示该写入项不是 tensor 写入项，例如 WriteItemType.BYTE_IO，此时返回 None。 2. 当 WriteItem.tensor_data 不为 None 时，根据 tensor_data.size 计算元素数量，并结合 tensor_data.properties.dtype 计算单个元素字节数，最终返回 tensor 存储字节数。 3. 该 API 不执行实际 tensor 数值计算，不涉及 NPU kernel 或底层算子调用，主要依赖 checkpoint planner 中已有的 tensor 元数据信息。 二、用例补齐说明 经检索，当前测试目录中未发现直接调用 WriteItem.tensor_storage_size() 并断言其返回值的专项测试用例。已有 checkpoint planner 相关测试主要覆盖 WriteItem、SavePlan、TensorWriteData 等结构和 planner 流程，但没有单独验证 tensor_storage_size 的返回值行为。 因此，本 PR 新增测试文件： test/distributed/checkpoint/test_planner_api.py 该测试文件覆盖以下场景： WriteItemType.TENSOR + tensor_data 构造普通 tensor 写入项，覆盖 torch.float32、torch.float16 和 torch.int8 三种 dtype，验证返回值等于： tensor.numel() * tensor.element_size() WriteItemType.SHARD + tensor_data 构造 shard 写入项，验证 tensor_storage_size 能够根据 shard 对应的 tensor metadata 返回正确的存储字节数。 WriteItemType.BYTE_IO 且 tensor_data 为 None 构造非 tensor 写入项，验证 tensor_storage_size 返回 None。 三、NPU 适配说明 torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 本身是元数据辅助方法，不直接依赖 NPU 算子。但测试中需要基于 tensor 构造 TensorWriteData 和 TensorProperties，因此测试 tensor 需要运行在 NPU 设备上。 本 PR 使用如下方式获取当前 accelerator： device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" 并通过如下方式将测试 tensor 迁移到当前设备： tensor = torch.empty((2, 3), dtype=dtype).to(device_type) 因此，在 NPU 环境下，测试 tensor 会运行在 NPU 设备上，满足 NPU 适配要求。 四、涉及版本说明 本次用例补齐分别提交到以下目标分支： v2.7.1 v2.9.0 v2.10.0 v2.11.0 v2.12.0 master 各版本均单独提交 PR，目标分支与源码分支一一对应。 # 【资料变更】 > 请确认是否涉及资料变更。如涉及，需要在PR中体现，并简要说明修改内容。如不涉及，需填写“不涉及” 涉及资料支持情况检查。本 PR 不直接修改资料文件。 已检查 v2.7.1 分支 docs/zh/native_apis，当前文档中已有 torch.distributed.checkpoint.planner.WriteItem，但未检索到 torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 的方法级支持信息。 已单独提交资料补齐 PR 到 v2.7.1 分支，补充 docs/zh/native_apis 中该 API 的支持情况。资料 PR 刷新以下版本目录： 1. docs/zh/native_apis/pytorch_2-7-1/torch-distributed-checkpoint.md 2. docs/zh/native_apis/pytorch_2-9-0/torch-distributed-checkpoint.md 3. docs/zh/native_apis/pytorch_2-10-0/torch-distributed-checkpoint.md pytorch_2-6-0 和 pytorch_2-8-0 按要求不刷新。 资料补齐 PR：<https://gitcode.com/Ascend/pytorch/pull/35673> # 【接口变更】 > 请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及，需要详细说明接口以及对应的变更内容，同时需要在资料中体现。如不涉及，需填写“不涉及” 不涉及 # 【功能验证】 > 说明测试场景，测试方法。如果本次测试方式与常规单元测试不同，请详细说明您的测试步骤\ > 新增/变更内容是否已新增/适配UT测试用例看护，并补充测试自验证截图 新增测试用例： test/distributed/checkpoint/test_planner_api.py 测试命令： python test/distributed/checkpoint/test_planner_api.py -v 测试场景： WriteItemType.TENSOR + tensor_data，验证返回 tensor 存储字节数。 WriteItemType.SHARD + tensor_data，验证返回 shard tensor 存储字节数。 WriteItemType.BYTE_IO 且 tensor_data 为 None，验证返回 None。 已在以下版本执行该用例： v2.7.1：通过 v2.9.0：通过 v2.10.0：通过 v2.11.0：通过 v2.12.0：通过 master：通过 不同版本运行用例结果示例： V2.7.1 [root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.7.1/test/distributed/checkpoint/test_planner_api.py -v test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok ---------------------------------------------------------------------- Ran 3 tests in 4.995s OK [root@734774725e04 workspace]# V2.9.0 [root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.9.0/test/distributed/checkpoint/test_planner_api.py -v test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok ---------------------------------------------------------------------- Ran 3 tests in 1.023s OK [root@734774725e04 workspace]# V2.10.0 [root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.10.0/test/distributed/checkpoint/test_planner_api.py -v test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok ---------------------------------------------------------------------- Ran 3 tests in 1.036s OK [root@734774725e04 workspace]# V2.11.0 [root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.11.0/test/distributed/checkpoint/test_planner_api.py -v test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok ---------------------------------------------------------------------- Ran 3 tests in 1.021s OK [root@734774725e04 workspace]# V2.12.0 [root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.12.0/test/distributed/checkpoint/test_planner_api.py -v test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok ---------------------------------------------------------------------- Ran 3 tests in 1.010s OK [root@734774725e04 workspace]# master [root@734774725e04 workspace]# python /workspace/ascend-pytorch-master/test/distributed/checkpoint/test_planner_api.py -v test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok ---------------------------------------------------------------------- Ran 3 tests in 1.026s OK [root@734774725e04 workspace]# 【CheckList】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [x] 代码注释完备，正确记录错误日志 - [x] 代码实现进行了返回值、空指针等校验 - [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [x] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!35505	3 个月前
test_staging_api.py	test: add staging api npu tests Co-authored-by: zj<2943168611@qq.com> # message auto-generated for no-merge-commit merge: !35469 merge test/2.9.0-staging-api-npu into v2.9.0 test: add staging api npu tests Created-by: zjucn Commit-by: zj Merged-by: ascend-robot Description: <!-- PR描述模板更新日期：20260203 --> # 【合入来源】 - [x] issue/工单: https://gitcode.com/Ascend/pytorch/issues/1897 本 PR 覆盖以下 API： - torch.distributed.checkpoint.staging.AsyncStager - torch.distributed.checkpoint.staging.AsyncStager.should_synchronize_after_execute - torch.distributed.checkpoint.staging.AsyncStager.stage - torch.distributed.checkpoint.staging.AsyncStager.synchronize_staging - torch.distributed.checkpoint.staging.BlockingAsyncStager - torch.distributed.checkpoint.staging.BlockingAsyncStager.stage - torch.distributed.checkpoint.staging.BlockingAsyncStager.synchronize_staging 问题背景和 API 功能说明见 issue。本 PR 不修改 torch-npu 生产代码，仅新增 test_staging_api.py，用于验证 DCP staging API 在 NPU tensor 场景下的可用性。 # 【修改方案】 ## 一、总体方案 PyTorch 社区已有分布式 checkpoint 测试主要通过 torch.distributed.checkpoint.async_save 端到端流程间接覆盖 staging 行为。本 PR 在 torch-npu 侧补充 staging API 专项测试，重点验证 NPU tensor 经过 DCP staging 后能够形成稳定的 CPU staged state dict，并可继续通过 DCP 文件系统 writer 完成异步保存和加载。 测试实现遵循以下原则： - 使用真实 NPU tensor 作为输入，不 mock DCP staging 内部逻辑。 - 直接调用 BlockingAsyncStager.stage()，验证 NPU 到 CPU staging 行为。 - 使用同一个 BlockingAsyncStager 实例连续 staging，验证 cached staging buffer 复用。 - 使用 FileSystemWriter.stage() 覆盖文件系统 writer 继承 BlockingAsyncStager 的真实路径。 - 使用 torch.distributed.checkpoint.async_save 和 torch.distributed.checkpoint.load 覆盖保存加载端到端链路。 - 使用自定义 writer 验证 should_synchronize_after_execute 与 synchronize_staging() 的调用关系。 - 用例保持轻量，不引入分布式初始化和多进程依赖，使测试目标聚焦在本 PR 涉及的 staging API。 ## 二、用例覆盖 本 PR 新增的测试文件为 test_staging_api.py。用例覆盖如下： ### 1. test_async_stager_protocol_defaults 覆盖 torch.distributed.checkpoint.staging.AsyncStager、torch.distributed.checkpoint.staging.AsyncStager.should_synchronize_after_execute、torch.distributed.checkpoint.staging.AsyncStager.stage 和 torch.distributed.checkpoint.staging.AsyncStager.synchronize_staging。 该用例构造最小 stager 实现，验证： - 对象可被 runtime protocol 识别为 AsyncStager。 - 默认 should_synchronize_after_execute 返回 True。 - 默认 stage() 抛出 NotImplementedError。 - 默认 synchronize_staging() 返回 None。 PyTorch 2.9.0 起 AsyncStager protocol 新增 close()，测试中的最小实现类提供该方法，以兼容不同版本的 runtime protocol 判断。 ### 2. test_blocking_async_stager_stage_npu_tensor_to_cpu_snapshot 覆盖 torch.distributed.checkpoint.staging.BlockingAsyncStager.stage。 该用例构造包含顶层和嵌套 NPU tensor 的 state dict，调用 BlockingAsyncStager.stage() 后验证： - staged tensor 位于 CPU。 - staged tensor 数值与 staging 时刻的源 NPU tensor 一致。 - 修改原始 NPU tensor 后，staged CPU tensor 不受影响。 该用例验证 BlockingAsyncStager.stage() 在 NPU 输入下的 CPU staging 和 checkpoint 快照语义。 ### 3. test_blocking_async_stager_cached_reuses_cpu_buffer 覆盖 torch.distributed.checkpoint.staging.BlockingAsyncStager.stage(cache_staged_state_dict=True)。 该用例使用同一个 BlockingAsyncStager 实例连续 staging 两份 NPU tensor，验证： - 两次返回的 staged tensor 复用同一 CPU buffer。 - 第二次 staging 后，CPU buffer 中的数据被更新为最新 NPU tensor 数据。 该用例验证 cached staging buffer 在 NPU 到 CPU 拷贝场景下的正确性。 ### 4. test_blocking_async_stager_sync_noop_and_property 覆盖 torch.distributed.checkpoint.staging.BlockingAsyncStager、torch.distributed.checkpoint.staging.AsyncStager.should_synchronize_after_execute 和 torch.distributed.checkpoint.staging.BlockingAsyncStager.synchronize_staging。 该用例验证： - BlockingAsyncStager 满足 AsyncStager 协议。 - BlockingAsyncStager.should_synchronize_after_execute 返回 False。 - BlockingAsyncStager.synchronize_staging() 为 no-op，并返回 None。 该用例验证阻塞式 staging 实现与 AsyncStager 协议的同步语义一致。 ### 5. test_filesystem_writer_stage_sets_copy_ahead_zero 覆盖 FileSystemWriter.stage() 与 BlockingAsyncStager.stage() 的组合路径。 该用例显式设置 per_thread_copy_ahead=1024 后调用 FileSystemWriter.stage()，验证： - per_thread_copy_ahead 被置为 0。 - 返回的 staged tensor 位于 CPU。 - staged tensor 数值与输入 NPU tensor 一致。 该用例验证文件系统 writer 在 async staging 场景下使用 CPU staged 数据路径。 ### 6. test_async_save_npu_state_dict_with_filesystem_writer 覆盖 torch.distributed.checkpoint.async_save、FileSystemWriter、BlockingAsyncStager.stage() 和 torch.distributed.checkpoint.load 的完整链路。 该用例分别测试 cache_staged_state_dict=False 和 cache_staged_state_dict=True，验证： - NPU tensor 可通过 DCP async_save 异步保存。 - 保存完成后可通过 DCP load 加载回 NPU 目标 tensor。 - 加载后的 tensor 数值与源 NPU tensor 一致。 该用例验证默认文件系统 writer 下的真实 DCP NPU 保存加载路径。 ### 7. test_async_save_converts_stateful_object 覆盖 torch.distributed.checkpoint.async_save 在 staging 前对 Stateful 对象的转换路径。 该用例定义实现 state_dict() 和 load_state_dict() 的对象，对象内部持有 NPU tensor，验证： - Stateful 对象可直接传给 async_save。 - 对象内部 NPU tensor 可被 staging 并保存。 - 加载后可恢复到新的对象中，且数值一致。 该用例验证 Stateful 对象与 DCP staging 调用链在 NPU tensor 场景下可用。 ### 8. test_async_save_calls_synchronize_when_stager_requests_it 覆盖 torch.distributed.checkpoint.staging.AsyncStager.should_synchronize_after_execute 与 torch.distributed.checkpoint.staging.AsyncStager.synchronize_staging 的调用关系。 该用例自定义 FileSystemWriter，将 _synchronize_after_execute 设为 True，并记录 synchronize_staging() 是否被调用。测试验证： - async_save 启动后台保存任务后会调用 synchronize_staging()。 - 保存 future 可正常完成。 该用例验证自定义 stager 请求同步时，async_save 能正确执行同步 hook。 ## 三、NPU 适配说明 本 PR 的 NPU 适配重点是验证 DCP staging API 在真实 NPU tensor 输入下的行为。 具体适配点如下： - 使用 torch.arange(...).to(device_type)、torch.ones(...).to(device_type) 和 torch.full(...).to(device_type) 构造 NPU tensor。 - 直接验证 NPU tensor 可通过 BlockingAsyncStager.stage() staging 到 CPU。 - 验证 staged CPU tensor 与原始 NPU tensor 解耦，避免后续训练修改影响 checkpoint 数据。 - 验证 cached staging buffer 可在多次 NPU 到 CPU staging 中复用并正确更新。 - 验证 FileSystemWriter.stage() 会在 async staging 场景下关闭 per_thread_copy_ahead。 - 验证 DCP.async_save + FileSystemWriter + DCP.load 能保存并恢复 NPU tensor。 - 验证 Stateful 对象内部 NPU tensor 可经 DCP staging 调用链保存和恢复。 为什么不需要修改 API： - AsyncStager 和 BlockingAsyncStager 是 DCP Python 层协议和默认实现，接口本身不绑定具体设备后端。 - NPU 适配风险集中在 NPU tensor 到 CPU staging、cached buffer 复用以及 FileSystemWriter 与 async_save 调用链是否可用。 - 本 PR 通过新增测试覆盖这些风险点，无需修改 torch-npu 生产代码或 API 签名。 ## 四、跨版本兼容说明 经对比 PyTorch 2.7.1、2.9.0、2.10.0、2.11.0 以及 master，本 PR 覆盖的七个 API 核心行为保持兼容，无需新增额外测试分支。 版本差异和本 PR 处理如下： - PyTorch 2.7.1 中，AsyncStager.stage 返回普通 STATE_DICT_TYPE。本 PR 测试覆盖该返回形式。 - PyTorch 2.9.0 起，AsyncStager.stage 返回类型扩展为普通 state dict 或 future；BlockingAsyncStager.stage 仍返回普通 staged state dict，本 PR 测试无需变化。 - PyTorch 2.9.0 起，AsyncStager.synchronize_staging 被标记为 deprecated，但该方法仍存在并保持同步 hook 语义。本 PR 继续验证该 API 的现有行为。 - PyTorch 2.9.0 起，AsyncStager protocol 新增 close()。测试中的最小 stager 实现提供 close()，保证 runtime protocol 判断在不同版本上兼容。 - PyTorch 2.9.0 起新增 DefaultStager、StagingOptions、AsyncSaveResponse 等 API；这些 API 不属于本 PR 覆盖范围，因此本 PR 不新增对应测试。 - PyTorch 2.10.0、2.11.0、2.12.0 和 master 相比 2.9.0，在本 PR 覆盖的七个 API 上未发现需要额外处理的公共语义变化。 # 【资料变更】 > 不涉及 # 【接口变更】 > 不涉及 # 【功能验证】 > 在 2.7.1、2.9.0、2.10.0、2.11.0、2.12.0 版本上执行 test_staging_api.py 测试，均通过。 bash $ python test_staging_api.py ........ ---------------------------------------------------------------------- Ran 8 tests in 5.361s OK # 【CheckList】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [x] 代码注释完备，正确记录错误日志 - [x] 代码实现进行了返回值、空指针等校验 - [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [x] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!35469	3 个月前