已合并
test: adapt Generator.clone_state test for NPU #43344
test: adapt Generator.clone_state test for NPU #43344
已合并
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.7.1 官方社区中已存在 torch.Generator.clone_state 相关用例,检索结果包括:

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

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

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

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

2. v2.7.1原有patch情况

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

原始官方用例中使用了 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,而 PyTorch 内部 common_cuda.py 中存在与 tuple 比较的逻辑,可能导致导入阶段触发异常:

TypeError: '>=' not supported between instances of 'NoneType' and 'tuple'

3. 本次修改内容

本次在已有 test_upstream/test/test_cuda.py.patch 基础上补充 test_graphsafe_set_get_rng_state 的 NPU 适配。

主要修改如下:

  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. 增加 torch.cuda.get_device_capability() 返回 None 时的兼容处理,避免测试导入阶段触发 NoneTypetuple 比较异常:
if torch.cuda.get_device_capability() is None:
    torch.cuda.get_device_capability = lambda device=None: (0, 0)  # noqa: E731
  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.7.1+cpu
torch-npu版本:2.7.1

将最终 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

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

【CheckList】

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

变更摘要

此 PR 将 test_cuda.py 补丁文件中的 test_graphsafe_set_get_rng_state 测试用例适配至 NPU 后端,并同步完善了 CUDA/NPU 双后端兼容的基础设施。核心改动包括:将测试中所有 CUDA 特定 API(如 CUDAGraphStreamdefault_generatorscurrent_stream 等)替换为 NPU 对应接口,同时在模块导入处新增 torch_npu.contrib.transfer_to_npu 导入、HAS_NPU 标志及 get_device_capability 回退逻辑,并在多处 memory_stats 相关调用中增加了 NPU 分支判断。

主要改动

  • NPU 基础兼容设施:在模块导入部分新增 import torch_npu.contrib.transfer_to_npuHAS_NPU 变量定义,以及 torch.cuda.get_device_capability 在 NPU 环境下的兜底 lambda,确保后续 NPU 检测与降级逻辑可用
  • memory_stats 系列调用双后端适配:将 test_cuda.py 中所有 torch.cuda.memory_stats()memory_stats_as_nested_dict()reset_peak_memory_stats()reset_accumulated_memory_stats()empty_cache() 调用替换为条件表达式,当 HAS_NPU 为真时使用 torch.npu.* 对应接口
  • test_graphsafe_set_get_rng_state 完整 NPU 适配:将该测试用例中的所有设备标识从 "cuda" 改为 "npu"torch.cuda.CUDAGraph 替换为 torch.npu.NPUGraphtorch.cuda.Stream 替换为 torch.npu.Streamtorch.cuda.default_generators 替换为 torch.npu.default_generators,CUDA 流操作(streamcurrent_stream)均切换为 NPU 流 API,并新增 @skipCUDAMemoryLeakCheckIf(True)@skipCUDANonDefaultStreamIf(True) 装饰器
  • test_multinomial_invalid_probs_cuda 设备选择适配:在大张量分配处增加 device="cuda" if not HAS_NPU else device="npu" 条件判断,使测试在 NPU 环境下正确定向到 NPU 设备
  • test_cublas_workspace_explicit_allocation NPU 跳过:在该测试函数上新增 @unittest.skipIf(HAS_NPU, "skip test on npu") 装饰器,避免 NPU 环境下执行不兼容的 cuBLAS 相关测试
likedislike
不准确?
atomgit-bot
atomgit-bot
7月30日 评论:

代码审查

经过全面审查,我确认除了已报告的 P2 问题外,diff 中没有其他需要报告的问题。多轮扫描的要点如下:

  • multinomial 批处理逻辑(第 92-103 行):batch_size=1000while remaining > 0 循环、min(remaining, batch_size) 边界处理均正确,torch.cat 拼接结果与原单次调用语义等价(replacement=True 下每次抽取独立)。
  • get_device_capability monkey-patch(第 14-15 行):仅在 CUDA 不可用时生效,属于 NPU 环境的防御性规避措施,不影响 CUDA 系统。
  • TEST_NPU = True(第 24 行):硬编码值,仅用于跳过 NPU 不支持的 cublas 测试,与整个 patch 文件依赖 torch_npu 可导入的设计一致。
  • @skipCUDAMemoryLeakCheckIf(True) / @skipCUDANonDefaultStreamIf(True):正确禁用了与 NPU 测试无关的 CUDA 特定检查。
  • torch.rand(1, device="npu") 初始化(第 156 行):在访问 default_generators[0] 前初始化 NPU 上下文,防御性但无害(随后 manual_seed(0) 会重置状态)。

审查总结

优先级 数量
P0 0
P1 0
P2 1
P3 0

已审查文件:

  • test_upstream/test/test_cuda.py.patch — 发现 1 个问题(P2:test_graphsafe_set_get_rng_state 缺少 NPU graph 能力检查)

整体风险评估:低风险。 此变更是一个标准的 NPU 适配 patch,将 CUDA 测试用例转换为 NPU 等价版本。唯一值得关注的问题是新增的 test_graphsafe_set_get_rng_state 测试的 skip 条件未检查 NPU graph 能力,可能在无 graph 支持的 NPU 环境中导致测试崩溃而非优雅跳过。其余变更均为正常的 API 替换(torch.cuda.*torch_npu.npu.* / torch.npu.*)和 NPU 限制的规避处理(multinomial 批处理),逻辑正确且与仓库中其他 patch 文件风格一致。

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

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

/approve

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