| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
test(distributed/checkpoint): Add StorageWriter storage_meta API tests Co-authored-by: Jwerr<19166910919@163.com> # message auto-generated for no-merge-commit merge: !35323 merge test-StorageWriter-api-master into master 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版本上执行该用例，均通过，本分支对应master日志如下： I0511 13:42:19.353000 130119 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:895] Started process 0 with pid 130263 I0511 13:42:19.355000 130119 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:895] Started process 1 with pid 130264 I0511 13:42:30.185000 130263 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:973] Process 0 skipping test test_storage_meta_checkpoint_id_consistency_across_ranks for following reason: Multi-NPU 2 condition not satisfied I0511 13:42:30.322000 130264 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:973] Process 1 skipping test test_storage_meta_checkpoint_id_consistency_across_ranks for following reason: Multi-NPU 2 condition not satisfied sI0511 13:42:33.087000 130119 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:895] Started process 0 with pid 130483 I0511 13:42:33.088000 130119 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:895] Started process 1 with pid 130484 I0511 13:42:43.797000 130484 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:973] Process 1 skipping test test_storage_meta_returned_in_distributed_save for following reason: Multi-NPU 2 condition not satisfied I0511 13:42:43.938000 130483 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:973] Process 0 skipping test test_storage_meta_returned_in_distributed_save for following reason: Multi-NPU 2 condition not satisfied sI0511 13:42:46.624000 130119 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:895] Started process 0 with pid 130675 I0511 13:42:46.625000 130119 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:895] Started process 1 with pid 130676 I0511 13:42:57.504000 130675 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:973] Process 0 skipping test test_storage_meta_save_id_changes_after_reset_distributed for following reason: Multi-NPU 2 condition not satisfied I0511 13:42:57.625000 130676 pytorch-2.12.0-env/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py:973] Process 1 skipping test test_storage_meta_save_id_changes_after_reset_distributed for following reason: Multi-NPU 2 condition not satisfied s.../root/work/pytorch-2.12.0-env/lib/python3.10/site-packages/torch/distributed/checkpoint/utils.py:485: UserWarning: torch.distributed is disabled, unavailable or uninitialized, assuming the intent is to save in a single process. return func(*args, **kwargs) .[W511 13:43:04.696461912 NPUCachingAllocator.cpp:199] Warning: The current CANN and Soc versions require processing for 32 padding size, with memory allocation. (function operator()) ... ---------------------------------------------------------------------- Ran 10 tests in 45.440s OK (skipped=3) # 【CheckList】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [x] 代码注释完备，正确记录错误日志 - [x] 代码实现进行了返回值、空指针等校验 - [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [x] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!35323 | 4 个月前 | |
test(distributed):add test for validate_checkpoint_id, reset, set_up_storage_reader from BroadcastingTorchSaveReader for master(fixed) Co-authored-by: xh-zhan<zhanxuhan@sjtu.edu.cn> # message auto-generated for no-merge-commit merge: !33996 merge test-master into master test(distributed):add test for validate_checkpoint_id, reset, set_up_storage_reader from BroadcastingTorchSaveReader for master(fixed) 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) ### Related PR This PR is a follow-up fix to the previously merged [PR #33733](https://gitcode.com/Ascend/pytorch/pull/33733). ### 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合入master 分支。 5. **备注**：本PR在原有基础上，将文件路径改为相对路径，并确保tensor在npu环境下执行。 See merge request: Ascend/pytorch!33996 | 5 个月前 | |
test: add NPU tests for BroadcastingTorchSaveReader APIs Co-authored-by: zj<2943168611@qq.com> # message auto-generated for no-merge-commit merge: !35060 merge test/master-broadcasting-reader-api-npu into master 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!35060 | 4 个月前 | |
| 2 年前 | ||
| 2 年前 | ||
| 2 年前 | ||
test: add load planner api npu tests Co-authored-by: zj<2943168611@qq.com> # message auto-generated for no-merge-commit merge: !35073 merge test/master-loadplan-api-npu into master 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!35073 | 4 个月前 |
| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 4 个月前 | ||
| 5 个月前 | ||
| 4 个月前 | ||
| 2 年前 | ||
| 2 年前 | ||
| 2 年前 | ||
| 4 个月前 |