已关闭

【API一致性任务】: torch.Generator.clone_state NPU一致性验证 #3476

Flipped创建于 7月30日关闭于 23 天前

7月30日创建

ISSUE来源：【社区任务】7月社区任务第三期-Ascend for PyTorch API 一致性开发（124） #3235

环境信息

操作系统：Ubuntu 22.04.5 LTS
机器架构：aarch64
昇腾硬件信息：Ascend 910B3 NPU
npu-smi版本：25.2.0
CANN软件版本：9.1.0-beta.1

安装的对应软件版本：

- torch 2.7.1+cpu，torch-npu 2.7.1
- torch 2.9.0+cpu，torch-npu 2.9.0.post2
- torch 2.10.0+cpu，torch-npu 2.10.0
- torch 2.11.0+cpu，torch-npu 2.11.0rc1
- torch 2.12.0+cu130，torch-npu 2.12.0rc1

使用场景及问题

当前需确认以下 API 在 NPU 环境下的测试覆盖、适配情况和资料支持情况：

torch.Generator.clone_state

经检查，PyTorch 官方已有直接测试用例：

test/test_cuda.py::TestCuda.test_graphsafe_set_get_rng_state

该用例在 NPU 上需要将 CUDA Generator、CUDA Graph、CUDA Stream 等相关逻辑适配为 NPU 对应接口，因此本任务走 1.1 路线：PyTorch 官方有用例，并且需要进行 NPU 适配。

按照当前版本要求，本次测试 patch 处理：

v2.7.1
v2.11.0
v2.12.0

v2.9.0 和 v2.10.0 已进入维护阶段，暂不提交测试 patch；master 不需要提交 test_upstream patch。

资料补齐已转移到 master 分支下，因此资料部分需要检查：

pytorch_2-7-1
pytorch_2-9-0
pytorch_2-10-0
pytorch_2-11-0
pytorch_2-12-0

PyTorch 2.6 和 2.8 按当前规则无需刷新。

一、API功能说明

torch.Generator.clone_state 用于克隆当前 Generator 的随机数状态，返回一个新的 Generator 对象。

克隆后的 Generator 应保持与原 Generator 克隆时一致的 RNG 状态，从而保证后续随机数生成结果可复现。

基础行为包括：

原 Generator 设置固定随机种子后，调用 clone_state() 得到新的 Generator；
使用原 Generator 和克隆 Generator 分别生成随机数；
两者从同一 RNG 状态开始，生成结果应保持一致；
克隆 Generator 的设备应与原 Generator 保持一致。

NPU 基础功能验证结果如下：

[PASS] torch.Generator.clone_state basic NPU behavior
generator device: npu
clone device: npu

因此，该 API 在 NPU 上实际已支持，本次无需修改 torch.Generator.clone_state 接口实现，只需补齐官方测试用例的 NPU 适配。

二、官方用例覆盖情况

PyTorch 官方中检索到 torch.Generator.clone_state 相关用例。

PyTorch v2.7.1

test/test_cuda.py:1965
test/test_cuda.py:2068
test/inductor/test_cudagraph_trees.py:2673

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

test/test_cuda.py::TestCuda.test_graphsafe_set_get_rng_state
test/test_cuda.py:L1951-L2028

PyTorch v2.11.0

test/test_cuda.py:2152
test/test_cuda.py:2255
test/inductor/test_cudagraph_trees.py:4822
test/test_xpu.py:1256

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

test/test_cuda.py::TestCuda.test_graphsafe_set_get_rng_state
test/test_cuda.py:L2138-L2215

PyTorch v2.12.0

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 状态一致性

因此，无需在 torch-npu 的 test/ 目录新增自定义测试用例，只需适配现有 PyTorch 官方用例。

三、NPU适配分析

官方目标用例原始逻辑使用 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 上的行为。

原始运行结果表现为：

test_graphsafe_set_get_rng_state (__main__.TestCuda.test_graphsafe_set_get_rng_state) ... skipped 'CUDA not available, skipping tests'

----------------------------------------------------------------------
Ran 1 test in 0.001s

OK (skipped=1)

因此，本次需要将该官方 CUDA 用例适配为 NPU 用例，使其在 NPU 环境中真实执行，而不是跳过。

主要适配点如下：

增加 NPU 相关导入：

import torch_npu
from torch_npu.contrib import transfer_to_npu

增加 torch.cuda.get_device_capability() 返回 None 时的兼容处理，避免测试导入阶段受到 CUDA capability 判断逻辑影响：

if torch.cuda.get_device_capability() is None:
    torch.cuda.get_device_capability = lambda device=None: (0, 0)  # noqa: E731

将目标用例中的 CUDA 设备张量替换为 NPU 设备张量：

torch.rand(..., device="npu")

将 Generator 初始化设备从 CUDA 替换为 NPU：

generator = torch.Generator(device="npu")

将 CUDA Graph / Stream 替换为 NPU Graph / Stream：

g = torch.npu.NPUGraph()
s = torch.npu.Stream()
default_generator = torch.npu.default_generators[0]

将 CUDA stream 上下文替换为 NPU stream 上下文：

with torch.npu.stream(s):

将 CUDA stream 同步替换为 NPU stream 同步：

torch.npu.current_stream().wait_stream(s)

将原 CUDA graph 可用性判断替换为 NPU 可用性判断：

@unittest.skipIf(
    not torch.npu.is_available(), "NPU not available, skipping tests"
)

增加以下装饰器，避免 CUDA memory leak / non-default stream 检查逻辑影响 NPU 目标用例运行：

@skipCUDAMemoryLeakCheckIf(True)
@skipCUDANonDefaultStreamIf(True)

修改后，该用例能够在 NPU 上真实执行并覆盖 torch.Generator.clone_state 的核心行为。

四、各版本patch处理

v2.7.1

v2.7.1 已存在：

test_upstream/test/test_cuda.py.patch

原 patch 已对 test/test_cuda.py 中部分 CUDA 相关测试进行了 NPU 适配，但未覆盖 torch.Generator.clone_state 对应的官方目标用例：

test/test_cuda.py::TestCuda.test_graphsafe_set_get_rng_state

本次在已有 test_upstream/test/test_cuda.py.patch 基础上补充该目标用例的 NPU 适配。

对应处理路线为 1.1.2：修改已有 patch。

v2.11.0

v2.11.0 原本不存在：

test_upstream/test/test_cuda.py.patch

本次新增：

test_upstream/test/test_cuda.py.patch

并完成 test/test_cuda.py::TestCuda.test_graphsafe_set_get_rng_state 的 NPU 适配。

对应处理路线为 1.1.1：新增 patch。

v2.12.0

v2.12.0 原本不存在：

test_upstream/test/test_cuda.py.patch

本次新增：

test_upstream/test/test_cuda.py.patch

并完成 test/test_cuda.py::TestCuda.test_graphsafe_set_get_rng_state 的 NPU 适配。

v2.12.0 中 test_graphsafe_set_get_rng_state 后存在其他 CUDA graph 测试函数，例如：

test_graph_rng_after_failed_capture

因此本次生成 patch 时已检查修改范围，确保只修改 test_graphsafe_set_get_rng_state 目标函数，不误改后续其他 CUDA graph 测试函数。

最终 v2.12.0 patch 修改范围已收敛为：

1 file changed, 88 insertions(+)
create mode 100644 test_upstream/test/test_cuda.py.patch

对应处理路线为 1.1.1：新增 patch。

五、功能验证

最终 patch 分别应用到对应 PyTorch 干净工作树后，重新运行官方目标用例。

v2.7.1

运行命令：

cd /workspace

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

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

python /workspace/pytorch-v2.7.1/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.191s

OK

v2.11.0

运行命令：

cd /workspace

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

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

python /workspace/pytorch-v2.11.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.190s

OK

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

三个目标版本修改后的官方用例均在 NPU 环境下执行通过，不再是 skipped，能够正常覆盖：

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

六、资料检查

已检查 master 分支资料目录：

docs/zh/api/native_api

torch.Generator.clone_state 为 PyTorch Generator 公开接口，不属于 PyTorch 私有接口。该 API 用于克隆当前 Generator 的随机数状态，返回一个新的 Generator 对象；克隆后的 Generator 应保持与原 Generator 克隆时一致的 RNG 状态，用于保证后续随机数生成结果可复现。

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

最新资料 PR 已完成以下补充：

本次资料 PR 修改范围为 master 最新资料目录中实际存在版本对应的 torch.md 文件：

docs/zh/api/native_api/pytorch_2-7-1/torch.md
docs/zh/api/native_api/pytorch_2-11-0/torch.md
docs/zh/api/native_api/pytorch_2-12-0/torch.md
docs/zh/api/native_api/pytorch_2-13-0/torch.md

补充后的资料条目格式如下：

clone_state()

原生文档：torch.Generator.clone_state

是否支持：是

限制与说明：-

综上，torch.Generator.clone_state 需要进行资料补齐，已通过 master 分支资料 PR 在现有版本的 torch.md 中完成支持情况刷新。

七、PR信息

v2.7.1：https://gitcode.com/Ascend/pytorch/pull/43344
v2.11.0：https://gitcode.com/Ascend/pytorch/pull/43346
v2.12.0：https://gitcode.com/Ascend/pytorch/pull/43348
master 资料 PR：https://gitcode.com/Ascend/pytorch/pull/43350

八、结论

torch.Generator.clone_state 已有 PyTorch 官方直接测试用例覆盖。
该 API 在 NPU 上基础功能验证通过，不需要修改接口实现。
本任务满足 1.1 路线要求：PyTorch 官方有用例，并且需要进行 NPU 适配。
本次测试 patch 处理版本为 v2.7.1、v2.11.0、v2.12.0。
v2.7.1 修改已有 test_upstream/test/test_cuda.py.patch。
v2.11.0 和 v2.12.0 分别新增 test_upstream/test/test_cuda.py.patch。
三个版本的 patch 均从对应 PyTorch 干净工作树通过 git diff 生成。
三个目标版本的 test_graphsafe_set_get_rng_state 均重新验证通过，且结果为 ok，不是 skipped。
v2.12.0 已额外检查 patch 修改范围，确认未误改后续其他 CUDA graph 测试函数。
v2.9.0 和 v2.10.0 已进入维护阶段，暂不提交测试 patch。
master 不需要提交 test_upstream patch。
无需新增 torch-npu 自定义测试用例。
无需修改 torch.Generator.clone_state 接口实现。
资料检查发现 torch.Generator.clone_state 在 master 分支 docs/zh/api/native_api 下五个目标版本均缺失资料条目。
资料补齐需单独提交 master PR。
资料补齐应新增 torch-Generator.md，不应写入 torch-random.md。
torch.Generator.clone_state 为非计算类 API，与数据类型无关，资料中限制与说明填写为 -。