已合并
test: adapt Generator.clone_state test for NPU #43348
test: adapt Generator.clone_state test for NPU #43348
已合并
Flipped创建于 7月30日
Flipped
Flipped
7月30日

【合入来源】

issue:【社区任务】7月社区任务第三期-Ascend for PyTorch API 一致性开发(124) #3235

issue:[Usage]: torch.Generator.clone_state NPU一致性验证 #3476

【修改方案】

1. API基础功能及用例覆盖

torch.Generator.clone_state 用于克隆当前 Generator 的随机数状态,返回一个新的 Generator 对象。克隆后的 Generator 应保持与原 Generator 克隆时一致的 RNG 状态,从而保证后续随机数生成结果可复现。

PyTorch v2.12.0 官方社区中已存在 torch.Generator.clone_state 相关用例,检索结果包括:

test/dynamo/test_after_aot.py:139
test/dynamo/test_after_aot.py:156
test/inductor/test_cudagraph_trees.py:5436
test/test_cuda.py:2249
test/test_cuda.py:2447
test/test_xpu.py:1319

其中本次适配的官方目标用例为:

test/test_cuda.py::TestCuda.test_graphsafe_set_get_rng_state
test/test_cuda.py:L2237-L2321

该用例直接覆盖以下 Generator RNG 状态相关功能:

  • graphsafe_get_state
  • clone_state
  • graphsafe_set_state
  • 默认 Generator 在图场景下的 RNG 状态一致性

因此,本次无需新增自定义测试用例,只需适配现有 PyTorch 官方用例。

2. v2.12.0原有patch情况

v2.12.0 原本不存在:

test_upstream/test/test_cuda.py.patch

因此,本次为 v2.12.0 新增:

test_upstream/test/test_cuda.py.patch

原始官方用例中使用了 CUDA 专属接口和 CUDA 设备:

torch.Generator(device="cuda")
torch.cuda.CUDAGraph()
torch.cuda.Stream()
torch.cuda.default_generators[0]
torch.cuda.stream(s)
torch.cuda.current_stream().wait_stream(s)
torch.rand(..., device="cuda")

在 NPU 环境下直接运行时,该用例会因 CUDA graph 条件不满足而跳过,无法覆盖 torch.Generator.clone_state 在 NPU 上的行为。

此外,在导入测试依赖时,torch.cuda.get_device_capability() 在 NPU 环境下可能返回 None 或触发兼容性告警,因此本次增加兼容处理,避免测试导入阶段受到 CUDA capability 判断逻辑影响。

3. 本次修改内容

本次新增 test_upstream/test/test_cuda.py.patch,对 test_graphsafe_set_get_rng_state 进行 NPU 适配。

主要修改如下:

  1. 增加 NPU 相关导入:
import torch_npu
from torch_npu.contrib import transfer_to_npu
  1. 增加 torch.cuda.get_device_capability() 返回 None 时的兼容处理:
if torch.cuda.get_device_capability() is None:
    torch.cuda.get_device_capability = lambda device=None: (0, 0)  # noqa: E731
  1. 将目标用例中的 CUDA 设备张量替换为 NPU 设备张量:
torch.rand(..., device="npu")
  1. 将 Generator 初始化设备从 CUDA 替换为 NPU:
generator = torch.Generator(device="npu")
  1. 将 CUDA graph 相关接口替换为 NPU graph 对应接口:
g = torch.npu.NPUGraph()
s = torch.npu.Stream()
default_generator = torch.npu.default_generators[0]
  1. 将 CUDA stream 上下文替换为 NPU stream 上下文:
with torch.npu.stream(s):
  1. 将 CUDA stream 同步替换为 NPU stream 同步:
torch.npu.current_stream().wait_stream(s)
  1. 将原 CUDA graph 可用性判断替换为 NPU 可用性判断:
@unittest.skipIf(
    not torch.npu.is_available(), "NPU not available, skipping tests"
)
  1. 增加以下装饰器,避免 CUDA memory leak / non-default stream 检查逻辑影响 NPU 目标用例运行:
@skipCUDAMemoryLeakCheckIf(True)
@skipCUDANonDefaultStreamIf(True)

导入顺序说明

test/test_cuda.pytorch_npu / transfer_to_npu 的导入位置属于初始化依赖例外。

TestCuda 类存在类级跳过条件:

@unittest.skipIf(not TEST_CUDA, "CUDA not available, skipping tests")
class TestCuda(TestCase):

TEST_CUDA 来自 torch.testing._internal.common_cuda,会在模块导入阶段计算。若严格将 torch_npu / transfer_to_npu 放到所有 torch.testing 导入之后,NPU 环境下 TEST_CUDA 会提前计算为 False,导致目标用例 test_graphsafe_set_get_rng_state 被类级 skip,无法覆盖 torch.Generator.clone_state。

因此本 patch 保留 transfer_to_npu 在 common_cuda 导入前初始化。该处理不使用 try-except,不修改 TestCuda 类级 skip 条件,也不新增 NPU 专用测试类,仅用于保证官方目标用例能在 NPU 上真实执行。

本次只适配官方已有用例,不修改 torch.Generator.clone_state 接口实现。
本次适配的 test_graphsafe_set_get_rng_state 已能覆盖 torch.Generator.clone_state 在 NPU Generator RNG 状态克隆场景下的核心行为。

【资料变更】

已检查 master 分支 docs/zh/api/native_api 资料目录,torch.Generator.clone_state 为 PyTorch Generator 公开接口,非私有接口。该 API 属于 Generator 状态类 API,非计算类 API,与具体数据类型无关,限制与说明填写为 -

根据 master 分支最新资料目录规则,pytorch_2-9-0pytorch_2-10-0 目录已删除,本次不再补充;新增 pytorch_2-13-0 目录需要同步检查并刷新。资料补齐以 master 分支 docs/zh/api/native_api 当前实际存在的版本目录为准。

资料补齐已在 master 分支单独提交 PR。本次未新增独立的 torch-Generator.md 文件,而是在各版本已有 torch.md## Generators 章节中补充 clone_state() 条目,位置放在 set_state() 之后,保持 torch.Generator 相关 API 统一维护在 torch.md 中。

【接口变更】

不涉及。

本次仅适配 PyTorch 官方测试用例,不修改客户可见接口,不修改 torch.Generator.clone_state 功能实现。

【功能验证】

运行环境

操作系统:Ubuntu 22.04.5 LTS
机器架构:aarch64
昇腾硬件:Ascend 910B3 NPU
npu-smi版本:25.2.0
CANN版本:9.1.0-beta.1
Python版本:3.11
torch版本:2.12.0+cu130
torch-npu版本:2.12.0rc1

将最终 patch 应用到干净的 PyTorch v2.12.0 工作树后,运行官方目标用例。

cd /workspace

source /workspace/envs/torch-npu-2.12.0/bin/activate
source /usr/local/Ascend/ascend-toolkit/set_env.sh

unset PYTHONPATH
export PYTHONPATH=/workspace/pytorch-v2.12.0/test
export PYTHONNOUSERSITE=1
export PYTORCH_PRINT_REPRO_ON_FAILURE=0

python /workspace/pytorch-v2.12.0/test/test_cuda.py -v -k test_graphsafe_set_get_rng_state

验证结果:

test_graphsafe_set_get_rng_state (__main__.TestCuda.test_graphsafe_set_get_rng_state) ... ok

----------------------------------------------------------------------
Ran 1 test in 0.192s

OK

最终 patch 可以在干净工作树中正常应用,修改后的官方用例已在 NPU 上通过,可以覆盖 torch.Generator.clone_state 的核心功能。

最终提交内容如下:

[test-generator-clone-state-2.12.0 3cc500ca3] test: adapt Generator.clone_state test for NPU
 1 file changed, 88 insertions(+)
 create mode 100644 test_upstream/test/test_cuda.py.patch

【CheckList】

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 Flipped 的贡献)
FlippedFlipped
7月30日 创建了 pull request,commit 431f79bc
atomgit-bot
atomgit-bot
7月30日 评论:

变更摘要

本次 PR 为 test_upstream/test/test_cuda.py.patch 新增适配补丁,将 PyTorch 官方用例 test_graphsafe_set_get_rng_state 从 CUDA 后端迁移至 NPU(昇腾)后端。该用例覆盖 torch.Generatorgraphsafe_get_stateclone_stategraphsafe_set_state 等 RNG 状态相关接口在 NPU 图模式下的行为一致性,确保 NPU 设备上克隆后的 Generator 与原 Generator 保持相同的随机数状态。

主要改动

  • 新增 NPU 导入与兼容性处理:添加 import torch_npufrom torch_npu.contrib import transfer_to_npu,并对 torch.cuda.get_device_capability 做空值兜底(返回 (0, 0)),避免 CUDA 设备检测缺失时引发异常。
  • 替换测试跳过条件:将 @unittest.skipIf 的条件从 not TEST_CUDA_GRAPH 改为 not torch.npu.is_available(),使测试仅在有可用 NPU 设备时执行。
  • 屏蔽 CUDA 专用检查装饰器:新增 @skipCUDAMemoryLeakCheckIf(True)@skipCUDANonDefaultStreamIf(True) 装饰器,关闭测试中对 CUDA 内存泄漏和非默认流的检查。
  • 设备与 API 全面 NPU 化:将所有 device="cuda" 替换为 device="npu",所有 CUDA 专属类(torch.cuda.CUDAGraphtorch.cuda.Stream)和接口(torch.cuda.streamtorch.cuda.current_streamtorch.cuda.default_generators)替换为对应的 NPU 等价物(torch.npu.NPUGraphtorch.npu.Stream 等)。
  • 初始化默认 NPU Generator:在访问 torch.npu.default_generators[0] 之前插入 torch.rand(1, device="npu") 调用,确保默认 NPU Generator 在使用前被正确初始化。
likedislike
atomgit-bot
atomgit-bot
7月30日 评论:

代码审查

审查结束 — 总结

变更文件审查结果

文件 审查结果
test_upstream/test/test_cuda.py.patch 发现 1 个问题(P2)

问题统计

  • P0: 0
  • P1: 0
  • P2: 1(get_device_capability() 条件检测模式不够健壮,与同类 patch 不一致,在特定环境下可能导致测试模块导入失败)
  • P3: 0

整体风险评估

该 patch 的核心逻辑(将 CUDA API 替换为 NPU API、适配 graphsafe_set_get_rng_state 用例)是正确的,设备替换全面且一致。唯一的风险点在于模块级别的 get_device_capability() 条件检测采用了与仓库中其他同类 patch(test_nn.py.patchtest_jit.py.patch)不一致的脆弱模式——在条件中直接调用可能抛异常的函数,而非使用 torch.accelerator.current_accelerator() 检测或无条件 patch。若 transfer_to_npu 已预先 patch 了该函数则实际不会触发,但该模式缺乏防御性,建议统一为安全模式。整体风险:低。

⚠️ 已识别出整体风险,但无法提取行内评论,请参考整体评估。

likedislike
ascend-robotascend-robot成员
7月30日 添加了label:ascend-cla/yes
此处折叠了54条消息 查看更多
li_jing_hw成员
23 天前 评论:

/approve

likedislike
ascend-robotascend-robot成员
23 天前 添加了label:approvedlgtm
ascend-robotascend-robot成员
23 天前 删除了label:ci-pipeline-passed
ascend-robotascend-robot成员
23 天前 合入了pull request
ascend-robot
ascend-robot成员
23 天前 评论:
流水线 pytorch_gitcode_PR_multiVersion#13941 [ commitID:6050a2df ] 运行失败
likedislike