【合入来源】

• 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 测试，均通过。

$ python test_staging_api.py
........
----------------------------------------------------------------------
Ran 8 tests in 5.361s

OK

【CheckList】

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

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常