已合并
refactor(npu): route RNG generator lookup through the accelerator entry point as upstream does #44830
refactor(npu): route RNG generator lookup through the accelerator entry point as upstream does #44830
已合并
dwoai22创建于 8月18日
dwoai22
dwoai22
8月18日

【合入来源】

关联 issue:https://gitcode.com/Ascend/pytorch/issues/4148

Fixes #4148

设计文档:random_patch_reduction_rfc.md(Phase 1)
依赖:#44025(已合入)。

【修改方案】

将 Generator 的获取入口由 NPU 专属的 Python 元组收敛到 accelerator 框架的 C++ 统一
入口,为后续将 torch.npu.* 随机数 API 委托到 torch.accelerator.* 做准备。

  1. torch_npu/npu/random.py:7 处 Generator 获取由
    torch_npu.npu.default_generators[idx] 改为
    torch._C._accelerator_getDefaultGenerator(idx),覆盖 get_rng_state /
    set_rng_state / manual_seed / manual_seed_all / seed / seed_all /
    initial_seed。

  2. torch_npu/npu/__init__.py:_get_generator() 同样改走 accelerator 入口。
    _set_rng_state_offset / _get_rng_state_offset 经由该函数一并收敛。

  3. 移除 random.py 中因此变为未使用的 import torch_npu。

  4. 延迟初始化语义保持不变

  5. 对外 API 的签名、默认参数、docstring 均未改动。

【资料变更】

不涉及。

【接口变更】

不涉及。

【功能验证】

已在 NPU 环境完成以下验证,全部通过。

核心 RNG 与 accelerator 入口

pytest ./npu/test_npu.py ./npu/test_random_sampling.py \
  -k "test_manual_seed or test_torch_manual_seed_seeds_npu_devices or test_get_set_rng_state or test_get_set_rng_state_input_device or test_accelerator_get_default_generator or test_seed"

图模式 RNG

pytest ./npu/test_aclgraph_rng_state.py ./npu/test_aclnn_random_aclgraph.py

其他已执行的验证

  • 延迟初始化:全新进程中调用 manual_seed / manual_seed_all / seed /
    seed_all / _set_rng_state_offset,断言调用后 is_initialized() 仍为 False,
    再触发初始化并校验种子生效。确认 _lazy_call 的语义未因
    _accelerator_getDefaultGenerator 的 eager 初始化行为而改变。
  • 回归看护:test_transfer_to_npu.py::test_update_cuda_default_generators /
    test_generator,确认 default_generators 仍正常填充。
  • 多卡:test_npu.py::test_accelerator_get_default_generator_per_device、
    test_npu_multinpu.py::test_get_set_rng_state_all / test_rng_state_offset。
  • 集成:test_jit_trace.py、test_jit_script.py、test_utils.py。

test_npu.PNG
test_graph.PNG

【CheckList】

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 dwoai22 的贡献)
dwoai22dwoai22
8月18日 创建了 pull request,commit 121fe40a
dwoai22dwoai22
8月18日 关联了issue:[Refactor] npu.random APIs do not use upstream's '_accelerator_getDefaultGenerator'
atomgit-bot
atomgit-bot
8月18日 评论:

变更摘要

本 PR 为重构优化(关联 issue #4148,Phase 1),将 NPU 随机数生成器(Generator)的获取入口从 NPU 专属的 Python 元组 torch_npu.npu.default_generators[idx] 收敛到 accelerator 框架的 C++ 统一入口 torch._C._accelerator_getDefaultGenerator(idx),与上游实现保持一致,为后续将 torch.npu.* 随机数 API 委托到 torch.accelerator.* 做准备。该改动仅涉及内部获取路径的替换,对外 API 的签名、默认参数与 docstring 均未改动,延迟初始化语义保持不变。

主要改动

  • 统一 Generator 获取入口:torch_npu/npu/random.py 中 7 处 Generator 获取由 torch_npu.npu.default_generators[idx] 改为 torch._C._accelerator_getDefaultGenerator(idx),覆盖 get_rng_state、set_rng_state、manual_seed、manual_seed_all、seed、seed_all、initial_seed 等核心 RNG 接口。
  • 收敛 _get_generator() 路径:torch_npu/npu/__init__.py 中 _get_generator() 同样改走 accelerator 入口,使 _set_rng_state_offset、_get_rng_state_offset 经由该函数一并收敛。
  • 清理无用导入:移除 torch_npu/npu/random.py 中因上述改动而不再使用的 import torch_npu,并将 typing 导入缩减为仅保留 Union。
likedislike
不准确?
atomgit-bot
atomgit-bot
8月18日 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
此处折叠了176条消息 查看更多
ascend-robotascend-robot成员
28 天前 关闭了关联的issue
ascend-robotascend-robot成员
28 天前 合入了pull request
ascend-robot
ascend-robot成员
28 天前 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike
AtlasAccount
AtlasAccount成员
28 天前 评论:
流水线 pytorch_gitcode_PR_multiVersion#14707 [ commitID:1f323f67 ] 已完成
likedislike
xuyun15成员16 天前进行代码检视1
torch_npu/npu/random.py
@@ -29,3 +25,1 @@
29- if idx is None:
30- idx = current_device()
31- default_generator = torch_npu.npu.default_generators[idx]
25+ idx = _get_device_index(device, optional=True)
xuyun1516 天前评论:

兼容性有问题,这里修改可能会导致

megatron/core/tensor_parallel/random.py", line 53, in _get_cuda_rng_state
[rank4]: return torch.cuda.random.get_rng_state(device=device)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.11/site-packages/torch_npu/npu/random.py", line 25, in get_rng_state
[rank4]: idx = _get_device_index(device, optional=True)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/usr/local/lib/python3.11/site-packages/torch/accelerator/_utils.py", line 16, in _get_device_index
[rank4]: raise ValueError(
[rank4]: ValueError: cuda doesn't match the current accelerator npu. when instantiating RowParallelLinear

likedislike