Star854
1343
代码
代码
Issues1607
Pull Requests2208
项目讨论
Wiki
分析
Star854
1343
已合并
test(jit): add ScriptModule API alignment test cases [v2.7.1] #37630
TensorLake创建于 6月4日
test(jit): add ScriptModule API alignment test cases [v2.7.1] #37630
已合并
TensorLake创建于 6月4日
从TensorLake/torch-nputest/script-module-npu-v2.7.1合入到Ascend/pytorchv2.7.1
T
TensorLake
6月4日

【合入来源】

Fork: TensorLake/torch-npu
分支: test/script-module-npu-v2.7.1 → Ascend/pytorch:v2.7.1

【修改方案】

本 PR 为 #1861 的交付。社区用例情况：已在 PyTorch 官方社区 pytorch/test/jit/ 下搜索 ScriptModule 相关测试，无针对这 19 个 API 的独立测试用例，因此自行编写测试用例。

一、API 功能说明及行为分类

torch.jit.ScriptModule 通过 torch.jit.script() 创建，返回 RecursiveScriptModule 实例。其方法可见性由双层机制控制：

  1. RecursiveScriptModule 自己定义的方法（如 save()、extra_repr()）
  2. _compiled_methods_allowlist 白名单中的 nn.Module 方法（如 train()、to()、state_dict()）
  3. 不满足以上两条的方法被 _make_fail 替换，调用即抛 RuntimeError("xxx is not supported on ScriptModules")

本测试使用 torch.jit.script() 作为 canonical 创建方式，正确验证 allowlist 机制的行为。

按实测行为将 19 个 API 分为 4 类：

第一类：allowlist 内正常工作（11 个）

API 来源 行为说明
train(mode) 白名单 设置训练模式，返回 self，递归传播到子模块
eval() 白名单 等价于 train(False)，返回 self
zero_grad(set_to_none) 白名单 清零所有参数梯度
float() 白名单 将参数/buffer 转为 float32，返回 self，递归传播
double() 白名单 将参数/buffer 转为 float64。NPU 不支持 float64，自动降级为 float32
to(dtype/device) 白名单 支持多参数形式（to(dtype)、to(device)、to(device,dtype)、to(dtype=...)、to(str_device)、to()），递归传播
type(dst_type) 白名单 转换 dtype，传入 int32 等非浮点类型抛 RuntimeError
state_dict(...) 白名单 返回 OrderedDict，支持 prefix=、destination=、keep_vars=
save(f, _extra_files) RSM 自定义 委托 C++ self._c.save()，无返回值。支持 _extra_files 附加文件
save_to_buffer() RSM 自定义 序列化到 bytes 并返回
extra_repr() RSM 自定义 返回 "original_name=..." 格式字符串

第二类：_make_fail 封杀（5 个，与 CPU/GPU 行为一致）

API 异常消息（assertRaisesRegex 匹配）
requires_grad_() "requires_grad_ is not supported on ScriptModules"
to_empty() "to_empty is not supported on ScriptModules"
xpu() "xpu is not supported on ScriptModules"
get_buffer() "get_buffer is not supported on ScriptModules"
set_submodule() "set_submodule is not supported on ScriptModules"

这 5 个 API 的行为与设备无关，属于 PyTorch 架构层面的设计决定。

第三类：torch-npu 拦截（3 个）

API CPU 行为 NPU 行为
share_memory() 正常工作（storage 变 shared） RuntimeError: "share_memory is not supported in npu"
register_module() RuntimeError（PyTorch 限制） RuntimeError（torch-npu 先拦截）
register_parameter() RuntimeError（PyTorch 限制） RuntimeError（torch-npu 先拦截）

注：share_memory 在 GPU/CUDA 上为 no-op（官方文档明确：This is a no-op for CUDA tensors），NPU 行为与 GPU 不一致。CPU 上正常。

第四类：PyTorch 内部未完成实现（1 个）

API 行为
set_extra_state() RuntimeError: "Reached a code path in Module.set_extra_state() that should never be called"

虽在白名单中，但走到未实现代码路径，属于上游 PyTorch 的 known issue。

二、测试文件说明

新增文件：test/jit/test_script_module.py，共 67 个测试用例，按功能语义分为 10 个测试类：

TestScriptModuleTrainEval（11 个） — 训练状态控制

  • test_train_default_is_training / test_train_set_true_explicit / test_train_set_false
  • test_eval_sets_training_false / test_train_returns_self / test_eval_returns_self
  • test_train_eval_roundtrip / test_train_on_npu / test_eval_on_npu
  • test_train_propagates_to_submodules / test_eval_propagates_to_submodules

TestScriptModuleZeroGrad（5 个） — 梯度管理

  • test_zero_grad_no_error / test_zero_grad_clears_grads / test_zero_grad_set_to_none
  • test_zero_grad_set_to_none_false / test_zero_grad_backward_chain_on_npu

TestScriptModuleTo（9 个） — 设备/dtype 转换

  • test_to_dtype / test_to_device / test_to_returns_self
  • test_to_device_and_dtype / test_to_dtype_keyword / test_to_string_device
  • test_to_no_args_returns_self / test_to_propagates_to_submodules / test_to_npu_and_dtype

TestScriptModuleFloatDouble（7 个） — float/double 转换

  • test_float_returns_self / test_float_converts_params / test_float_on_npu
  • test_float_propagates_to_submodules
  • test_double_returns_self / test_double_converts_params / test_double_on_npu_fallback_to_float32

TestScriptModuleType（4 个） — type 转换

  • test_type_float32 / test_type_float64 / test_type_on_npu / test_type_int32_raises

TestScriptModuleStateDict（7 个） — 序列化状态

  • test_state_dict_contains_params / test_state_dict_contains_buffers
  • test_state_dict_values_match / test_state_dict_on_npu
  • test_state_dict_with_prefix / test_state_dict_with_destination / test_state_dict_keep_vars

TestScriptModuleSave（6 个） — 模型保存

  • test_save_and_load / test_save_preserves_output / test_save_returns_none
  • test_save_on_npu / test_save_with_extra_files / test_save_to_buffer

TestScriptModuleExtraRepr（3 个） — 额外描述

  • test_extra_repr_returns_str / test_extra_repr_contains_original_name / test_extra_repr_on_npu

TestScriptModuleShareMemory（4 个） — 共享内存

  • test_share_memory_cpu_returns_self / test_share_memory_cpu_makes_shared
  • test_share_memory_cpu_idempotent / test_share_memory_on_npu_raises

TestScriptModuleMetadata（7 个） — 结构/元数据

  • test_register_module_raises_on_npu / test_register_parameter_raises_on_npu
  • test_set_submodule_raises / test_set_submodule_nested_raises
  • test_get_buffer_unsupported / test_get_buffer_unsupported_on_nested / test_get_buffer_unsupported_nonexistent

TestScriptModuleUnsupported（4 个） — PyTorch 设计限制

  • test_requires_grad_unsupported / test_to_empty_unsupported / test_xpu_unsupported
  • test_set_extra_state_raises

三、NPU 适配说明

  • 所有涉及张量的测试均在 NPU 设备上运行，使用 torch.accelerator.current_accelerator() 获取设备
  • NPU 不支持 float64，double()/to(torch.float64)/type(torch.float64) 自动降级为 float32
  • share_memory 在 NPU 上被拦截（GPU 为 no-op，行为不一致，已在测试中标注）
  • register_module/register_parameter 在 NPU 上被拦截（与 CPU RuntimeError 消息不同但结果一致）
  • _make_fail 封杀的 5 个 API 在所有设备上行为一致

【资料变更】

经检查 docs/zh/native_apis/ 下各版本路径（v2.7.1 / v2.9.0 / v2.10.0 / v2.11.0 / v2.12.0）：

  • 已有记录的 API（18/19）：register_module、register_parameter、requires_grad_、save、set_extra_state、share_memory、state_dict、to、to_empty、train、type、xpu、zero_grad、double、eval、extra_repr、float、get_buffer
  • 缺失的 API（1/19）：set_submodule 在 torch-jit.md 中未找到对应条目，需补充

【接口变更】

不涉及

【功能验证】

测试文件：test/jit/test_script_module.py
测试环境：torch 2.7.1 + torch_npu 2.7.1 + NPU 910B3（CANN 8.5.0）

test_extra_repr_contains_original_name ... ok
test_extra_repr_on_npu ... ok
test_extra_repr_returns_str ... ok
test_double_converts_params ... [W604 ...] Warning: Device do not support double...
ok
test_double_on_npu_fallback_to_float32 ... ok
test_double_returns_self ... ok
test_float_converts_params ... ok
test_float_on_npu ... ok
test_float_propagates_to_submodules ... ok
test_float_returns_self ... ok
test_get_buffer_unsupported ... ok
test_get_buffer_unsupported_nonexistent ... ok
test_get_buffer_unsupported_on_nested ... ok
test_register_module_raises_on_npu ... ok
test_register_parameter_raises_on_npu ... ok
test_set_submodule_nested_raises ... ok
test_set_submodule_raises ... ok
test_save_and_load ... ok
test_save_on_npu ... ok
test_save_preserves_output ... ok
test_save_returns_none ... ok
test_save_to_buffer ... ok
test_save_with_extra_files ... ok
test_share_memory_cpu_idempotent ... ok
test_share_memory_cpu_makes_shared ... ok
test_share_memory_cpu_returns_self ... ok
test_share_memory_on_npu_raises ... ok
test_state_dict_contains_buffers ... ok
test_state_dict_contains_params ... ok
test_state_dict_keep_vars ... ok
test_state_dict_on_npu ... ok
test_state_dict_values_match ... ok
test_state_dict_with_destination ... ok
test_state_dict_with_prefix ... ok
test_to_device ... ok
test_to_device_and_dtype ... ok
test_to_dtype ... ok
test_to_dtype_keyword ... ok
test_to_no_args_returns_self ... ok
test_to_npu_and_dtype ... ok
test_to_propagates_to_submodules ... ok
test_to_returns_self ... ok
test_to_string_device ... ok
test_eval_on_npu ... ok
test_eval_propagates_to_submodules ... ok
test_eval_returns_self ... ok
test_eval_sets_training_false ... ok
test_train_default_is_training ... ok
test_train_eval_roundtrip ... ok
test_train_on_npu ... ok
test_train_propagates_to_submodules ... ok
test_train_returns_self ... ok
test_train_set_false ... ok
test_train_set_true_explicit ... ok
test_type_float32 ... ok
test_type_float64 ... ok
test_type_int32_raises ... ok
test_type_on_npu ... ok
test_requires_grad_unsupported ... ok
test_set_extra_state_raises ... ok
test_to_empty_unsupported ... ok
test_xpu_unsupported ... ok
test_zero_grad_backward_chain_on_npu ... ok
test_zero_grad_clears_grads ... ok
test_zero_grad_no_error ... ok
test_zero_grad_set_to_none ... ok
test_zero_grad_set_to_none_false ... ok

----------------------------------------------------------------------
Ran 67 tests in 2.866s
OK

结果：67 passed，0 failed。

【CheckList】

likedislike
Pull Request已成功合入, 合并人@ascend-robot
（感谢 TensorLake 的贡献）
TTensorLake
6月4日 创建了 pull request，commit e0e15d5d
ascend-robot
ascend-robot成员
6月4日 评论：
ascend-robotascend-robot成员
6月4日 添加了label：ascend-cla/yes
ascend-robot
ascend-robot成员
6月4日 评论：

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here

PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
test ✅ 李伟, sunyu-xuan (2/2) ✅ 李伟 (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

TensorLake, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
ascend-robot
ascend-robot成员
6月4日 评论：

当前仓库存在以下 保护分支 ：

Protected Branch Version Release
master
v2.7.1
v2.12.0
v2.10.0
v2.9.0
v2.11.0
ci-test
sync-pr28113--to-v2.9.0

评论 /sync <branch1> <branch2> ... 可将当前 PR 修改同步到其它分支（创建同步 PR）：
a) 如果当前 PR 是 Open 状态，同步操作将延迟到 PR 被合并时执行
b) 如果当前 PR 已经 Merged，将立即执行同步操作

注意：

  1. /sync 命令可以指定同步到多个分支，仅最后一个 /sync 命令生效
  2. 如果创建的同步 PR 不正确，可通过向同步 PR 的源分支提交轻量级 PR 完善，或使用 /close 命令关闭
likedislike
此处折叠了71条消息 查看更多
S
sunyu-xuan成员
6月8日 评论：

/lgtm

likedislike
L
liwei386成员
6月9日 评论：

/approve

likedislike
ascend-robotascend-robot成员
6月9日 添加了label：approvedlgtm
ascend-robotascend-robot成员
6月9日 合入了pull request
ascend-robot
ascend-robot成员
6月9日 评论：
流水线 pytorch_gitcode_PR_multiVersion#10325 [ commitID：e341db0b ] 已完成
likedislike