| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
test(distributed/checkpoint): Add StorageWriter storage_meta API tests Co-authored-by: Jwerr<19166910919@163.com> # message auto-generated for no-merge-commit merge: !35277 merge test-StorageWriter-api-v2.12.0 into v2.12.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.12.0日志如下： 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!35277 | 3 个月前 | |
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 | 3 个月前 | |
test: add NPU tests for BroadcastingTorchSaveReader APIs Co-authored-by: zj<2943168611@qq.com> # message auto-generated for no-merge-commit merge: !35291 merge test/2.12.0-broadcasting-reader-api-npu into v2.12.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!35291 | 3 个月前 | |
| 2 年前 | ||
Add test for _EmptyStateDictLoadPlanner on NPU Co-authored-by: yuhuan2580<wn0010240@163.com> # message auto-generated for no-merge-commit merge: !42765 merge EmptyStateDictLoadPlanner-v2.12.0 into v2.12.0 Add test for _EmptyStateDictLoadPlanner on NPU Created-by: ljfnjj1234567 Commit-by: yuhuan2580 Merged-by: ascend-robot Description: <!-- PR描述模板更新日期：20260203 --> # 【合入来源】 关联任务issue: https://gitcode.com/Ascend/pytorch/issues/3146 - [ ] 需求 - [ ] 问题单 - [x] issue/工单 - [ ] 重构优化 - [ ] 资料更新 # 【修改方案】 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.12.0/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.12.0 configfile: pyproject.toml plugins: hypothesis-6.156.6 collected 15 items user_data/pytorch-v2.12.0/test/distributed/checkpoint/test_empty_state_dict_load_planner.py ............... [100%] ================================================ 15 passed in 4.13s ================================================ # 【CheckList】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [x] 代码注释完备，正确记录错误日志 - [x] 代码实现进行了返回值、空指针等校验 - [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [x] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!42765 | 4 天前 | |
| 2 年前 | ||
| 2 年前 | ||
test: add load planner api npu tests Co-authored-by: zj<2943168611@qq.com> # message auto-generated for no-merge-commit merge: !35294 merge test/2.12.0-loadplan-api-npu into v2.12.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!35294 | 3 个月前 | |
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: !35401 merge test-writeitem-tensor-storage-size-v2.12.0 into v2.12.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!35401 | 2 个月前 | |
test: add staging api npu tests Co-authored-by: zj<2943168611@qq.com> # message auto-generated for no-merge-commit merge: !35466 merge test/2.12.0-staging-api-npu into v2.12.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!35466 | 3 个月前 | |
[test] Add _save_state_dict Api Coverage Co-authored-by: lihaokun-2026<851563813@qq.com> # message auto-generated for no-merge-commit merge: !42814 merge save_state_dict-v2.12.0 into v2.12.0 [test] Add _save_state_dict Api Coverage Created-by: lihaokun-2026 Commit-by: lihaokun-2026 Merged-by: ascend-robot Description: <!-- PR描述模板更新日期：20260203 --> # 【合入来源】 [#3167](https://gitcode.com/Ascend/pytorch/issues/3167) - [ ] 需求 - [ ] 问题单 - [x] issue/工单 - [ ] 重构优化 - [ ] 资料更新 公开issue： - torch.distributed.checkpoint.state_dict_saver._save_state_dict: [【社区任务】7月社区任务第三期-Ascend for PyTorch API 一致性开发（59）](https://gitcode.com/Ascend/pytorch/issues/3167) # 【修改方案】 1. API 功能说明： - 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。 2. 上游用例检查情况： 检索 PyTorch / Torch-NPU 现有测试后，未找到面向 torch.distributed.checkpoint.state_dict_saver._save_state_dict 的独立、直接、最小验证用例，因此按任务要求新增 NPU API 验证用例。 3. 新增用例文件： - test_state_dict_saver__save_state_dict.py 4. 新增用例覆盖场景： | 测试场景 | 对应的测试方法 | 主要验证内容 | |---|---|---| | 普通 NPU Tensor 保存 | test_save_npu_tensor_no_dist | NPU Tensor 能否正常保存并生成元数据和检查点文件 | | 嵌套状态字典保存 | test_save_nested_npu_state_dict_no_dist | 嵌套字典是否能正确展开为扁平键并写入元数据 | | 自定义 SavePlanner | test_save_npu_tensor_with_custom_planner_no_dist | _save_state_dict 是否正确调用自定义 planner 的完整流程 | | 禁用 collectives | test_save_npu_tensor_with_collectives_disabled_no_dist | use_collectives=False 时能否正常保存 | # 【资料变更】 不涉及。未修改产品资料、接口说明文档或用户可见文档。 # 【接口变更】 不涉及。未修改跨代码仓接口，也未修改用户可见接口行为。 # 【功能验证】 测试环境： sh torch: 2.12.0+cpu torch_npu: 2.12.0 测试方法： sh 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】 > PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x] - [x] 代码注释完备，正确记录错误日志 - [x] 代码实现进行了返回值、空指针等校验 - [x] PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等 - [x] PR持续集成流水线（CI）执行通过，代码检查无异常 See merge request: Ascend/pytorch!42814 | 4 天前 | |
[v2.12.0] test: add test cases for _write_item Co-authored-by: huowentan<3294682143@qq.com> # message auto-generated for no-merge-commit merge: !42896 merge api-write-item-v2.12.0 into v2.12.0 [v2.12.0] test: add test cases for _write_item Created-by: 2501_93637465 Commit-by: huowentan Merged-by: ascend-robot Description: # 【合入来源】 任务来源：https://gitcode.com/Ascend/pytorch/issues/3148 关联 Issue：https://gitcode.com/Ascend/pytorch/issues/3148 Fixes #3148 - [ ] 需求 - [ ] 问题单 - [x] issue/工单 - [ ] 重构优化 - [ ] 资料更新 API： torch.distributed.checkpoint.filesystem._write_item 目标分支：v2.12.0 # 【修改方案】 经检索，PyTorch 社区现有用例缺少对 torch.distributed.checkpoint.filesystem._write_item 的独立、充分验证，因此新增测试文件： test/distributed/checkpoint/test_write_item.py 测试覆盖： - BytesIO 基本数据、空数据和大数据写入； - Tensor 基本写入； - float32、int64 和 bool Tensor； - 二维 Tensor 和标量 Tensor； - 多个写入项连续写入； - Tensor 与 BytesIO 混合连续写入； - stream 写入内容验证； - WriteResult 的 index、size、offset、length 和 storage key； - _StorageInfo 属性验证； - 不同 PyTorch 版本 _write_item 签名的兼容调用。 该 API 的 Tensor 写入路径要求输入数据位于 CPU，并通过 torch.save 写入文件流。传入 NPU Tensor 会违反该接口的 data.device == cpu 约束， 因此本测试按照接口定义使用 CPU Tensor，不进行 NPU Tensor 迁移。 # 【资料变更】 torch.distributed.checkpoint.filesystem._write_item 是 PyTorch 私有接口。 根据 API 一致性任务指南，私有接口无需补充 docs/zh/api/native_api 支持资料，因此本 PR 不涉及资料变更。 # 【接口变更】 不涉及接口定义或功能实现变更，仅新增 API 一致性测试。 # 【功能验证】 | 测试项 | 结果 | | --- | --- | | test/distributed/checkpoint/test_write_item.py | PASS | | 13 个测试用例 | PASS | | python -m py_compile | PASS | | git diff --check | PASS | | PR CI 流水线 | PASS | 测试结果： text ---------------------------------------------------------------------- Ran 13 tests OK # 【CheckList】 - [x] 使用完整 API 名称 - [x] 新增文件包含 Huawei copyright - [x] 新增文件包含模块 docstring - [x] 使用 TestCase、run_tests 和 self.assert* - [x] CPU Tensor 使用符合 _write_item 接口约束 - [x] 不包含 print、try-except 或 import torch_npu - [x] 当前分支仅包含一个提交 - [x] CLA 和 CI 检查通过 See merge request: Ascend/pytorch!42896 | 3 天前 |
| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 3 个月前 | ||
| 3 个月前 | ||
| 3 个月前 | ||
| 2 年前 | ||
| 4 天前 | ||
| 2 年前 | ||
| 2 年前 | ||
| 3 个月前 | ||
| 2 个月前 | ||
| 3 个月前 | ||
| 4 天前 | ||
| 3 天前 |