【合入来源】

https://gitcode.com/Ascend/pytorch/issues/2026
请勿携带内部流程信息（需求链接、问题单、内部issue等）

[x] issue/工单

【修改方案】

一、API功能说明

本 PR 补充 torch.utils.checkpoint 弱覆盖 API 的直接验证用例：
- torch.utils.checkpoint.SelectiveCheckpointContext：selective checkpoint policy function 接收的上下文对象，用于区分 forward/recompute 阶段；2.12.0 起还可携带 op_output。
- torch.utils.checkpoint.detach_variable：reentrant checkpoint 内部用于 detach 输入 Tensor 并保持 requires_grad 状态的辅助函数。

二、测试用例完备性说明

PyTorch 官方社区 test/test_autograd.py 已通过 checkpoint 调用链间接覆盖上述 API，但缺少直接、聚焦的 API 行为断言。本 PR 按 1.3 类补充 torch-npu 自定义测试文件 test/test_utils_checkpoint_api.py。

1. torch.utils.checkpoint.detach_variable 直接覆盖：
   - test_detach_variable_keeps_device_and_requires_grad 直接调用 torch.utils.checkpoint.detach_variable((x, y, marker))。
   - x 和 y 均为 NPU Tensor，覆盖返回 Tensor 保持 NPU device、requires_grad 与原输入一致、grad_fn 被断开并成为 leaf Tensor，非 Tensor 对象透传。
   - test_detach_variable_rejects_non_tuple_input 覆盖非 tuple 输入的 RuntimeError 分支。

2. torch.utils.checkpoint.SelectiveCheckpointContext 直接覆盖：
   - test_selective_checkpoint_context_direct_attributes 直接实例化 torch.utils.checkpoint.SelectiveCheckpointContext，断言对象类型和 is_recompute 属性；对支持 op_output 的版本，额外断言 op_output 指向 NPU Tensor 输出。
   - test_selective_checkpoint_context_passed_to_policy_fn 通过 create_selective_checkpoint_contexts + checkpoint(..., use_reentrant=False, context_fn=...) 捕获 policy_fn 收到的 ctx，断言 forward/recompute 两阶段均收到 SelectiveCheckpointContext，并验证上下文来自 NPU Tensor 的 selective checkpoint 运行路径。

上述新增用例把原来的调用链弱覆盖补充为直接 API 行为覆盖；其中涉及 Tensor 的路径均使用 torch.accelerator.current_accelerator() 获取 NPU 设备并在 NPU Tensor 上执行。

三、NPU适配方案

1. API 源码无需修改：torch.utils.checkpoint.SelectiveCheckpointContext 与 torch.utils.checkpoint.detach_variable 均复用 PyTorch Python 层实现，不新增 NPU kernel，不改变接口语义。
2. 测试新增路径：新增 test/test_utils_checkpoint_api.py，文件聚焦 checkpoint API 直接行为，可后续扩展同模块 API。
3. 分支提交范围：新增测试文件提交到 2.7.1、2.9.0、2.10.0、2.11.0、2.12.0 以及 master。

四、变更文件

• test/test_utils_checkpoint_api.py

【资料变更】

已检查资料支持情况，结论如下：

checkpoint：

• 公开 API，在文档中记录：torch.utils.checkpoint.CheckpointPolicy、torch.utils.checkpoint.SelectiveCheckpointContext、torch.utils.checkpoint.create_selective_checkpoint_contexts、torch.utils.checkpoint.set_checkpoint_debug_enabled。对应资料 PR：https://gitcode.com/Ascend/pytorch/merge_requests/36231
• 内部函数，不在文档中记录：torch.utils.checkpoint.detach_variable。

【接口变更】

不涉及。API 源码无修改，不新增、删除或变更对外接口。

【功能验证】

已在 2.10.0 对应 torch-npu 环境执行新增用例，结果如下；master 使用 2.12.0 验证环境执行。

验证环境

• 操作系统：Ubuntu / Linux aarch64
• 昇腾硬件：910B2
• CANN 软件版本：8.3.RC1
• 验证环境：
  • 2.7.1：torch 2.7.1+cpu / torch-npu 2.7.1.post3
  • 2.9.0：torch 2.9.0+cpu / torch-npu 2.9.0.post1
  • 2.10.0：torch 2.10.0+cpu / torch-npu 2.10.0rc3
  • 2.11.0：torch 2.11.0+cpu / torch-npu 2.11.0rc1
  • 2.12.0 / master 验证：torch 2.12.0+cpu / torch-npu 2.12.0.rc1

运行命令与结果

python test/test_utils_checkpoint_api.py
Ran 4 tests in 1.568s
OK

【CheckList】

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