Pull Request已成功合入, 合并人@ascend-robot
(感谢 dwoai22 的贡献)变更摘要
本 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。


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.




兼容性有问题,这里修改可能会导致
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


【合入来源】
关联 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.*做准备。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。torch_npu/npu/__init__.py:_get_generator()同样改走 accelerator 入口。_set_rng_state_offset/_get_rng_state_offset经由该函数一并收敛。移除
random.py中因此变为未使用的import torch_npu。延迟初始化语义保持不变
对外 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
其他已执行的验证
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。【CheckList】