【合入来源】

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

[x] issue/工单

【修改方案】

一、API功能说明

本 PR 覆盖两组 torch.utils API 的 NPU 测试适配：

1. torch.utils.checkpoint 相关 API
- torch.utils.checkpoint.CheckpointPolicy
- torch.utils.checkpoint.create_selective_checkpoint_contexts
- torch.utils.checkpoint.SelectiveCheckpointContext
- torch.utils.checkpoint.detach_variable
- torch.utils.checkpoint.set_checkpoint_debug_enabled

这些 API 用于 activation checkpoint 的 selective checkpoint 策略控制、forward/recompute 上下文创建、reentrant checkpoint 输入 detach 处理，以及 checkpoint debug 全局开关。

2. torch.utils.data DataLoader / collate / worker 相关 API
- torch.utils.data.DataLoader
- torch.utils.data._utils.collate.collate
- torch.utils.data.dataloader._BaseDataLoaderIter
- torch.utils.data.dataloader._SingleProcessDataLoaderIter
- torch.utils.data.dataloader._MultiProcessingDataLoaderIter
- torch.utils.data.dataloader._DatasetKind.create_fetcher
- torch.utils.data.dataloader._InfiniteConstantSampler
- torch.utils.data._utils.worker.WorkerInfo
- torch.utils.data._utils.worker.ManagerWatchdog
- torch.utils.data._utils.worker._IterableDatasetStopIteration
- torch.utils.data._utils.worker._ResumeIteration
- torch.utils.data._utils.signal_handling._set_worker_signal_handlers
- torch.utils.data._utils.signal_handling._remove_worker_pids

这些 API 覆盖 DataLoader 入口、单进程/多进程 iterator、Dataset fetcher 创建、IterableDataset 无限采样、batch collate、worker 信息、worker 存活检查、persistent worker 恢复、IterableDataset worker 结束通知以及 worker 信号处理和 pid 清理。

二、测试用例完备性说明

PyTorch 官方社区 test/test_autograd.py 和 test/test_dataloader.py 已有可复用用例。本 PR 不新增 torch-npu 自定义测试文件，只对官方用例中目标 API 相关 Tensor 做 NPU 设备适配，提交 test_upstream patch。

checkpoint API 覆盖情况：

1. torch.utils.checkpoint.set_checkpoint_debug_enabled
   覆盖类型：直接覆盖，强覆盖。
   覆盖用例：TestAutograd.test_checkpoint_detects_non_determinism。
   覆盖路径：测试代码显式执行 with torch.utils.checkpoint.set_checkpoint_debug_enabled(True/False)，再调用 checkpoint(..., debug=False/True)，验证全局 debug 上下文能覆盖 checkpoint 入参 debug，并验证非确定性重计算错误信息是否切换到 debug trace 分支。
   NPU 状态：参与 checkpoint forward/recompute 的输入 a 已迁移到 device="npu"，非确定性重计算检测在 NPU Tensor 上执行。

2. torch.utils.checkpoint.create_selective_checkpoint_contexts
   覆盖类型：直接覆盖，强覆盖。
   覆盖用例：TestSelectiveActivationCheckpoint.test_bad_inputs、test_output_already_has_autograd_meta、test_subclass_dispatching_sizes、test_policy_with_state、test_storage_lifetime、test_version_counter、test_function_with_more_than_one_output、test_function_with_non_tensor_output、test_can_only_trigger_recompute_once 等。
   覆盖路径：测试代码通过 create_selective_checkpoint_contexts 构造 context_fn，并传入 checkpoint(..., use_reentrant=False, context_fn=context_fn)，覆盖 OpOverload list、自定义 policy function、状态型 policy、allow_cache_entry_mutation、非法 op list、非法 OpOverloadPacket 和非法入参类型。
   NPU 状态：selective checkpoint 相关用例中的 x / y 等 Tensor 已迁移到 device="npu"，policy 触发的 forward、cache、recompute、backward 路径在 NPU Tensor 上执行。

3. torch.utils.checkpoint.CheckpointPolicy
   覆盖类型：直接覆盖，强覆盖。
   覆盖用例：TestSelectiveActivationCheckpoint 中多个 policy_fn 返回 CheckpointPolicy.MUST_SAVE / CheckpointPolicy.PREFER_RECOMPUTE 等枚举值。
   覆盖路径：policy function 返回 CheckpointPolicy 后，checkpoint 内部根据策略决定保存 forward 输出或在 backward 阶段重计算，覆盖保存、重计算、缓存 mutation 检测、版本计数、多输出、非 Tensor 输出和重复 backward 错误路径。
   NPU 状态：policy 作用的 aten op 输入/输出来自 NPU Tensor，缓存和重计算路径均在 NPU Tensor 上验证。

4. torch.utils.checkpoint.SelectiveCheckpointContext
   覆盖类型：调用链覆盖，当前为弱覆盖。
   覆盖用例：TestSelectiveActivationCheckpoint.test_policy_with_state 以及其他 selective checkpoint 用例。
   覆盖路径：create_selective_checkpoint_contexts 内部创建 SelectiveCheckpointContext，并在 _CachingTorchDispatchMode / _CachedTorchDispatchMode 中传入 policy function；test_policy_with_state 读取 ctx.is_recompute，验证 forward 与 recompute 阶段上下文状态切换。
   当前不足：测试没有直接实例化 SelectiveCheckpointContext，也没有对 context 对象类型进行显式断言；覆盖重点在 checkpoint 调用链中 ctx.is_recompute 的使用。
   NPU 状态：触发 ctx 的 policy 路径运行在 NPU Tensor 的 selective checkpoint 过程中。
   后续补用例：后续拆分补充 torch-npu test 目录用例，捕获 policy_fn 收到的 ctx，显式断言 isinstance(ctx, torch.utils.checkpoint.SelectiveCheckpointContext)、ctx.is_recompute 在 forward/recompute 阶段分别出现，并在支持 op_output 的版本中验证 ctx.op_output 与 NPU Tensor 输出的关系。

5. torch.utils.checkpoint.detach_variable
   覆盖类型：调用链覆盖，当前为弱覆盖。
   覆盖用例：TestAutograd.test_checkpoint_valid_reset_on_error。
   覆盖路径：测试执行 checkpoint(torch.exp, a, use_reentrant=True)，进入 reentrant checkpoint 的 CheckpointFunction.forward，源码中调用 detach_variable(tuple(inputs)) 对输入 Tensor detach 并保持 requires_grad 状态；用例同时验证异常后 checkpoint 状态恢复以及后续 backward 正常执行。
   当前不足：测试没有直接调用 torch.utils.checkpoint.detach_variable，也没有逐项断言 detach 后 Tensor 的 requires_grad、grad_fn、device 和非 Tensor passthrough 行为。
   NPU 状态：输入 a 已迁移到 device="npu"，detach_variable 处理的是 NPU Tensor。
   后续补用例：后续拆分补充 torch-npu test 目录用例，直接调用 detach_variable((x, y, marker))，断言返回 Tensor 保持 NPU device、requires_grad 与原输入一致、grad_fn 被断开、非 Tensor 对象保持透传，并覆盖非 tuple 输入的 RuntimeError 分支。

DataLoader / collate / worker API 覆盖情况：

1. torch.utils.data.DataLoader
   覆盖类型：直接覆盖，强覆盖。
   覆盖用例：TestDataLoader.test_sequential_batch、TestDataLoader.test_sequential_nonbatch、TestDataLoader.test_seqential_batch_workers、TestDataLoader.test_iterable_style_dataset、TestDataLoaderPersistentWorkers.test_dataset_not_reset 等。
   覆盖路径：测试直接构造 DataLoader，覆盖 batch_size=None、batch_size>0、num_workers=0、num_workers>0、IterableDataset、persistent_workers 等典型入口参数组合。
   NPU 状态：test_sequential_batch / test_sequential_nonbatch 中 TensorDataset 的 data / labels 已迁移到 NPU，验证 DataLoader 返回 NPU Tensor payload；worker 控制流用例主要验证多进程状态机，不强行把 worker 子进程内部数据迁移为 NPU Tensor。

2. torch.utils.data._utils.collate.collate
   覆盖类型：函数体覆盖，强覆盖；NPU Tensor 分支已覆盖。
   覆盖用例：TestDataLoader.test_default_collate_dtype、test_default_collate_mapping_keep_type、test_default_collate_sequence_keep_type、test_default_collate_bad_numpy_types、test_default_collate_shared_tensor。
   覆盖路径：用例调用 dataloader.default_collate，default_collate 内部进入 torch.utils.data._utils.collate.collate，覆盖 Tensor dtype、mapping 类型保持、sequence 类型保持、非法 numpy string/object 异常、shared tensor 分支。
   NPU 状态：test_default_collate_dtype 中补充 NPU Tensor batch，验证 collate 后输出仍位于 NPU；mapping/sequence/异常/shared-memory 分支保持官方原有 CPU/控制语义。

3. torch.utils.data.dataloader._BaseDataLoaderIter
   覆盖类型：调用链覆盖，强覆盖。
   覆盖用例：所有通过 iter(DataLoader) / for batch in DataLoader 执行的目标用例。
   覆盖路径：DataLoader.__iter__ 创建具体 iterator 后，执行 _BaseDataLoaderIter.__init__、__iter__、__next__、_next_index、_num_yielded 更新、IterableDataset len warning 相关公共流程。
   NPU 状态：单进程顺序迭代用例中该公共 next 流程处理 NPU Tensor payload。

4. torch.utils.data.dataloader._SingleProcessDataLoaderIter
   覆盖类型：调用链覆盖，强覆盖。
   覆盖用例：TestDataLoader.test_sequential_batch、TestDataLoader.test_sequential_nonbatch。
   覆盖路径：num_workers=0 时 DataLoader 创建 _SingleProcessDataLoaderIter，执行 _DatasetKind.create_fetcher、_next_index、fetch、collate 和可选 pin_memory 分支判断。
   NPU 状态：本 PR 将 TensorDataset 的 data / labels 迁移到 NPU，因此 _SingleProcessDataLoaderIter 读取和返回的是 NPU Tensor payload。

5. torch.utils.data.dataloader._MultiProcessingDataLoaderIter
   覆盖类型：调用链覆盖，当前为弱覆盖。
   覆盖用例：TestDataLoader.test_seqential_batch_workers、TestDataLoader.test_iterable_style_dataset、TestDataLoaderPersistentWorkers.test_dataset_not_reset。
   覆盖路径：num_workers>0 时创建 _MultiProcessingDataLoaderIter，覆盖 worker 启动、index_queue、worker_result_queue、prefetch、shutdown、persistent worker 复用等控制流。
   当前不足：当前用例验证的是多进程 DataLoader 控制面，未把 worker 子进程返回的数据强制改成 NPU Tensor；因此不能把它表述为 NPU Tensor payload 下的完整多进程 IPC 覆盖。
   NPU 状态：用例在 NPU 适配环境下运行；多进程控制 API 本身无 Tensor 计算语义。
   后续补用例：后续拆分补充更小的 torch-npu test 目录用例，围绕 num_workers=1、极小数据集、明确断言 iterator 类型、worker 生命周期和结果正确性；如确需覆盖 NPU Tensor 跨进程返回，则单独设计 NPU IPC 前置条件和资源约束，避免把 DataLoader 控制面问题与 NPU IPC 问题混在同一个验证点。

6. torch.utils.data.dataloader._DatasetKind.create_fetcher
   覆盖类型：调用链覆盖，强覆盖。
   覆盖用例：TestDataLoader.test_sequential_batch、test_sequential_nonbatch、test_iterable_style_dataset、test_seqential_batch_workers。
   覆盖路径：_SingleProcessDataLoaderIter.__init__ 和 worker_loop 均通过 _DatasetKind.create_fetcher 创建 MapDatasetFetcher 或 IterableDatasetFetcher；目标用例同时覆盖 map-style TensorDataset 和 IterableDataset。
   NPU 状态：map-style 单进程 fetcher 处理 NPU TensorDataset；IterableDataset/multi-worker 分支验证 fetcher 控制逻辑。

7. torch.utils.data.dataloader._InfiniteConstantSampler
   覆盖类型：调用链覆盖，当前为弱覆盖。
   覆盖用例：TestDataLoader.test_iterable_style_dataset。
   覆盖路径：IterableDataset 下 DataLoader 使用 _InfiniteConstantSampler 生成无限 None 索引，由 IterableDataset 自身控制 StopIteration。
   当前不足：用例通过 IterableDataset 行为间接证明 sampler 工作，没有直接实例化 _InfiniteConstantSampler 并断言连续 yield None。
   NPU 状态：该 API 不处理 Tensor，仅提供采样控制流。
   后续补用例：后续补充 torch-npu test 目录小用例，直接 iter(_InfiniteConstantSampler())，断言连续返回 None，并结合 IterableDataset DataLoader 验证不会提前终止。

8. torch.utils.data._utils.worker.WorkerInfo
   覆盖类型：直接覆盖，强覆盖。
   覆盖用例：TestDataLoader.test_get_worker_info。
   覆盖路径：子进程用例验证主进程 get_worker_info() 返回 None；worker_init_fn 和 worker dataset 内部 get_worker_info() 返回 WorkerInfo，并校验 id、num_workers、seed、dataset 字段和只读行为。
   NPU 状态：WorkerInfo 是 worker 元信息对象，不处理 Tensor；用例在 NPU 适配环境下执行。

9. torch.utils.data._utils.worker.ManagerWatchdog
   覆盖类型：调用链覆盖，当前为弱覆盖。
   覆盖用例：TestDataLoader.test_seqential_batch_workers、TestDataLoader.test_iterable_style_dataset、TestDataLoaderPersistentWorkers.test_dataset_not_reset。
   覆盖路径：worker_loop 中创建 ManagerWatchdog，并在 while watchdog.is_alive() 循环中驱动 worker 持续读取 index_queue。
   当前不足：用例覆盖正常 worker loop，但没有直接断言 ManagerWatchdog.is_alive()，也没有模拟父进程死亡场景。
   NPU 状态：该 API 是进程存活检查逻辑，不处理 Tensor。
   后续补用例：后续补充 torch-npu test 目录小用例，直接构造 ManagerWatchdog 并断言正常父进程场景 is_alive() 为 True；如需要父进程死亡场景，使用独立 subprocess 隔离验证。

10. torch.utils.data._utils.worker._IterableDatasetStopIteration
    覆盖类型：调用链覆盖，当前为弱覆盖。
    覆盖用例：TestDataLoader.test_iterable_style_dataset。
    覆盖路径：IterableDataset 多 worker 某个 worker 先耗尽时，worker_loop 捕获 StopIteration 并向主进程发送 _IterableDatasetStopIteration(worker_id)，主进程据此停止继续向该 worker 分发任务。
    当前不足：用例通过最终数据量、worker 正常退出和 warning 行为验证该控制消息链路，没有直接断言 _IterableDatasetStopIteration 对象字段。
    NPU 状态：该 API 是 dataclass 控制消息，不处理 Tensor。
    后续补用例：后续补充 torch-npu test 目录小用例，直接构造 _IterableDatasetStopIteration(worker_id) 并断言字段，同时保留 IterableDataset 多 worker 耗尽集成验证。

11. torch.utils.data._utils.worker._ResumeIteration
    覆盖类型：调用链覆盖，当前为弱覆盖。
    覆盖用例：TestDataLoaderPersistentWorkers.test_dataset_not_reset。
    覆盖路径：persistent_workers=True 跨 epoch 复用 worker 时，主进程向 worker 发送 _ResumeIteration，worker_loop 收到后重置 iteration_end 并重建 fetcher，保证 dataset 不重复初始化。
    当前不足：用例验证 persistent worker 复用效果，但没有直接断言 _ResumeIteration 对象字段或 ack 消息。
    NPU 状态：该 API 是 persistent worker 控制消息，不处理 Tensor。
    后续补用例：后续补充 torch-npu test 目录小用例，直接构造 _ResumeIteration(seed)，断言 seed 字段，并通过 persistent_workers=True 的最小 DataLoader 验证跨 epoch 复用行为。

12. torch.utils.data._utils.signal_handling._set_worker_signal_handlers
    覆盖类型：调用链覆盖，当前为弱覆盖。
    覆盖用例：TestDataLoader.test_seqential_batch_workers、TestDataLoader.test_iterable_style_dataset、TestDataLoaderPersistentWorkers.test_dataset_not_reset。
    覆盖路径：每个 DataLoader worker 进入 _utils.worker._worker_loop 后首先调用 signal_handling._set_worker_signal_handlers() 注册 SIGBUS、SIGSEGV、SIGFPE、SIGTERM 等 C 侧 handler。
    当前不足：用例覆盖 worker 正常启动路径，但没有直接制造 fatal signal 并检查 handler 输出。
    NPU 状态：该 API 是 C 侧 worker 信号处理入口，不处理 Tensor。
    后续补用例：后续补充 subprocess 隔离 smoke test，至少验证 worker 启动时该入口可被调用且不影响 DataLoader 正常退出；fatal signal 行为如需验证应单独隔离，避免影响主测试进程。

13. torch.utils.data._utils.signal_handling._remove_worker_pids
    覆盖类型：调用链覆盖，当前为弱覆盖。
    覆盖用例：TestDataLoader.test_seqential_batch_workers、TestDataLoader.test_iterable_style_dataset、TestDataLoaderPersistentWorkers.test_dataset_not_reset。
    覆盖路径：_MultiProcessingDataLoaderIter shutdown finally 分支在 _worker_pids_set 为 True 时调用 _utils.signal_handling._remove_worker_pids(id(self))，清理 C 侧 worker pid 记录。
    当前不足：用例验证 iterator 正常关闭和 worker exitcode，但没有直接检查 C 侧 pid registry 是否清空。
    NPU 状态：该 API 是 C 侧 pid registry 清理入口，不处理 Tensor。
    后续补用例：后续补充最小多进程 DataLoader 关闭用例，显式删除 iterator、join worker，并结合后续新 DataLoader 可正常运行来验证 pid 清理链路；不直接依赖 C 侧私有 registry 内部状态。

结论：当前 test_upstream patch 对 torch.utils.checkpoint.set_checkpoint_debug_enabled、torch.utils.checkpoint.create_selective_checkpoint_contexts、torch.utils.checkpoint.CheckpointPolicy、torch.utils.data.DataLoader、torch.utils.data._utils.collate.collate、torch.utils.data.dataloader._BaseDataLoaderIter、torch.utils.data.dataloader._SingleProcessDataLoaderIter、torch.utils.data.dataloader._DatasetKind.create_fetcher、torch.utils.data._utils.worker.WorkerInfo 属于强覆盖；对 torch.utils.checkpoint.SelectiveCheckpointContext、torch.utils.checkpoint.detach_variable、torch.utils.data.dataloader._MultiProcessingDataLoaderIter、torch.utils.data.dataloader._InfiniteConstantSampler、torch.utils.data._utils.worker.ManagerWatchdog、torch.utils.data._utils.worker._IterableDatasetStopIteration、torch.utils.data._utils.worker._ResumeIteration、torch.utils.data._utils.signal_handling._set_worker_signal_handlers、torch.utils.data._utils.signal_handling._remove_worker_pids 属于调用链弱覆盖，后续按 checkpoint 和 DataLoader 两组拆分补充更聚焦的 torch-npu 自定义用例。

三、NPU适配方案

1. API 源码无需修改：上述 API 属于 torch.utils.checkpoint 和 torch.utils.data 的 Python 控制面、autograd graph 管理、DataLoader 状态机和 collate 实现，不新增 NPU kernel，不改变对外接口。
2. checkpoint 适配方式：对 test/test_autograd.py 中目标 checkpoint 用例的默认 CPU Tensor 做最小 NPU 迁移，例如将参与 forward/recompute/backward 的输入改为 device="npu"。
3. DataLoader 适配方式：对 test/test_dataloader.py 中目标 DataLoader 单进程顺序迭代和 collate Tensor 分支做最小 NPU 迁移，将 TensorDataset data / labels 和 default_collate Tensor batch 放到 device="npu"。
4. 多进程说明：worker/signal/resume/stop-iteration API 的核心是 DataLoader 控制面。当前 PR 先保留官方多进程控制流验证，不把 worker 返回 Tensor 强行迁移为 NPU Tensor；后续弱覆盖补用例会单独拆分，避免将 DataLoader 控制面验证与 NPU IPC 验证混在同一个用例中。
5. 分支情况：2.12.0 基于当前最新 v2.12.0 分支重新生成 patch，避免携带旧目录和冲突内容。

四、变更文件

• test_upstream/test/test_autograd.py.patch
• test_upstream/test/test_dataloader.py.patch

【资料变更】

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

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。

dataloader：

• 公开 API，在文档中记录：torch.utils.data.DataLoader、torch.utils.data._utils.collate.collate。对应资料 PR：https://gitcode.com/Ascend/pytorch/merge_requests/36233
• 内部函数/类，不在文档中记录：DataLoader iterator、worker、signal handling、persistent worker 相关内部对象。

【接口变更】

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

【功能验证】

已在 v2.12.0 torch-npu 环境执行 NPU 适配后的 PyTorch 官方社区定向用例，结果如下：

验证环境

• 操作系统：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：torch 2.12.0+cpu / torch-npu 2.12.0.rc1

运行命令与结果

运行日志

# checkpoint
执行命令：python test/test_autograd.py TestSelectiveActivationCheckpoint TestAutograd.test_checkpoint_detects_non_determinism TestAutograd.test_checkpoint_valid_reset_on_error
Ran 13 tests in 1.111s
OK (skipped=1)

# DataLoader / collate / worker
执行命令：python test/test_dataloader.py TestDataLoader.test_sequential_batch TestDataLoader.test_sequential_nonbatch TestDataLoader.test_iterable_style_dataset TestDataLoader.test_seqential_batch_workers TestDataLoader.test_get_worker_info TestDataLoaderPersistentWorkers.test_dataset_not_reset TestDataLoader.test_default_collate_dtype TestDataLoader.test_default_collate_mapping_keep_type TestDataLoader.test_default_collate_sequence_keep_type TestDataLoader.test_default_collate_bad_numpy_types TestDataLoader.test_default_collate_shared_tensor
Ran 11 tests in 28.651s
OK

【CheckList】

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