test(jit): add ScriptModule API alignment test cases [v2.12.0]-pytorch-AtomGit

已合并

test(jit): add ScriptModule API alignment test cases [v2.12.0] #37637

TensorLake创建于 6月4日从TensorLake/torch-npu： test/script-module-npu-v2-12-0 合入到 Ascend/pytorch： v2.12.0

test(jit): add ScriptModule API alignment test cases [v2.12.0] #37637

已合并

TensorLake创建于 6月4日

从TensorLake/torch-npu test/script-module-npu-v2-12-0合入到Ascend/pytorch v2.12.0

TensorLake

6月4日

【合入来源】

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

【修改方案】

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

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

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

RecursiveScriptModule 自己定义的方法（如 save()、extra_repr()）
_compiled_methods_allowlist 白名单中的 nn.Module 方法（如 train()、to()、state_dict()）
不满足以上两条的方法被 _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)`	白名单	支持多参数形式，递归传播
`type(dst_type)`	白名单	转换 dtype
`state_dict(...)`	白名单	返回 OrderedDict，支持 `prefix=`、`destination=`、`keep_vars=` 参数
`save(f, _extra_files)`	RSM 自定义	委托 C++，无返回值。支持 `_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"`

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

API	CPU 行为	NPU 行为
`share_memory()`	正常工作	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（官方文档明确），NPU 行为与 GPU 不一致。

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

API	行为
`set_extra_state()`	RuntimeError: "should never be called"

二、测试文件说明

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

TestScriptModuleTrainEval（11 个）：train/eval 默认状态、显式设置、返回值验证、往返切换、NPU 验证、子模块递归传播
TestScriptModuleZeroGrad（5 个）：无梯度不报错、梯度清零、set_to_none 语义、set_to_none=False 清零、backward→zero_grad→backward 链路
TestScriptModuleTo（9 个）：dtype/device/device+dtype/dtype_keyword/string_device/no_args 多形式、返回值验证、子模块传播
TestScriptModuleFloatDouble（7 个）：float 转换/传播、double 转换/NPU float64 降级 float32/传播
TestScriptModuleType（4 个）：float32 正常路径、float64 降级路径、NPU 验证、int32 异常
TestScriptModuleStateDict（7 个）：参数/buffer 包含、值匹配、NPU 设备验证、prefix/destination/keep_vars 参数形式
TestScriptModuleSave（6 个）：save/load 往返、参数保留、返回 None、NPU 验证、_extra_files、save_to_buffer
TestScriptModuleExtraRepr（3 个）：返回字符串、original_name 模式、NPU 验证
TestScriptModuleShareMemory（4 个）：CPU 返回 self/makes_shared/幂等、NPU assertRaisesRegex 拦截
TestScriptModuleMetadata（7 个）：register_module/register_parameter NPU 拦截、set_submodule _make_fail、get_buffer _make_fail
TestScriptModuleUnsupported（4 个）：requires_grad_/to_empty/xpu _make_fail、set_extra_state 未完成实现

所有可预测异常的 API 均使用 assertRaisesRegex 精确匹配异常信息内容。

三、NPU 适配说明

所有涉及张量的测试均在 NPU 设备上运行
NPU 不支持 float64，double()/to(torch.float64)/type(torch.float64) 自动降级为 float32
share_memory 在 NPU 上被拦截（GPU 为 no-op，行为不一致，已在测试中标注）
_make_fail 封杀的 5 个 API 在所有设备上行为一致

【资料变更】

经检查 docs/zh/native_apis/ 下各版本路径，已有记录的 API 18/19。set_submodule 缺失，已通过独立文档 PR 提交。

【接口变更】

不涉及

【功能验证】

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

test_extra_repr_contains_original_name (__main__.TestScriptModuleExtraRepr.test_extra_repr_contains_original_name) ... ok
test_extra_repr_on_npu (__main__.TestScriptModuleExtraRepr.test_extra_repr_on_npu) ... ok
test_extra_repr_returns_str (__main__.TestScriptModuleExtraRepr.test_extra_repr_returns_str) ... ok
test_double_converts_params (__main__.TestScriptModuleFloatDouble.test_double_converts_params) ... [W604 11:34:32.815442779 compiler_depend.ts:164] Warning: Device do not support double dtype now, dtype cast replace with float. (function operator())
ok
test_double_on_npu_fallback_to_float32 (__main__.TestScriptModuleFloatDouble.test_double_on_npu_fallback_to_float32) ... ok
test_double_returns_self (__main__.TestScriptModuleFloatDouble.test_double_returns_self) ... ok
test_float_converts_params (__main__.TestScriptModuleFloatDouble.test_float_converts_params) ... ok
test_float_on_npu (__main__.TestScriptModuleFloatDouble.test_float_on_npu) ... ok
test_float_propagates_to_submodules (__main__.TestScriptModuleFloatDouble.test_float_propagates_to_submodules) ... ok
test_float_returns_self (__main__.TestScriptModuleFloatDouble.test_float_returns_self) ... ok
test_get_buffer_unsupported (__main__.TestScriptModuleMetadata.test_get_buffer_unsupported) ... ok
test_get_buffer_unsupported_nonexistent (__main__.TestScriptModuleMetadata.test_get_buffer_unsupported_nonexistent) ... ok
test_get_buffer_unsupported_on_nested (__main__.TestScriptModuleMetadata.test_get_buffer_unsupported_on_nested) ... ok
test_register_module_raises_on_npu (__main__.TestScriptModuleMetadata.test_register_module_raises_on_npu) ... ok
test_register_parameter_raises_on_npu (__main__.TestScriptModuleMetadata.test_register_parameter_raises_on_npu) ... ok
test_set_submodule_nested_raises (__main__.TestScriptModuleMetadata.test_set_submodule_nested_raises) ... ok
test_set_submodule_raises (__main__.TestScriptModuleMetadata.test_set_submodule_raises) ... ok
test_save_and_load (__main__.TestScriptModuleSave.test_save_and_load) ... ok
test_save_on_npu (__main__.TestScriptModuleSave.test_save_on_npu) ... ok
test_save_preserves_output (__main__.TestScriptModuleSave.test_save_preserves_output) ... ok
test_save_returns_none (__main__.TestScriptModuleSave.test_save_returns_none) ... ok
test_save_to_buffer (__main__.TestScriptModuleSave.test_save_to_buffer) ... ok
test_save_with_extra_files (__main__.TestScriptModuleSave.test_save_with_extra_files) ... ok
test_share_memory_cpu_idempotent (__main__.TestScriptModuleShareMemory.test_share_memory_cpu_idempotent) ... ok
test_share_memory_cpu_makes_shared (__main__.TestScriptModuleShareMemory.test_share_memory_cpu_makes_shared) ... ok
test_share_memory_cpu_returns_self (__main__.TestScriptModuleShareMemory.test_share_memory_cpu_returns_self) ... ok
test_share_memory_on_npu_raises (__main__.TestScriptModuleShareMemory.test_share_memory_on_npu_raises) ... ok
test_state_dict_contains_buffers (__main__.TestScriptModuleStateDict.test_state_dict_contains_buffers) ... ok
test_state_dict_contains_params (__main__.TestScriptModuleStateDict.test_state_dict_contains_params) ... ok
test_state_dict_keep_vars (__main__.TestScriptModuleStateDict.test_state_dict_keep_vars) ... ok
test_state_dict_on_npu (__main__.TestScriptModuleStateDict.test_state_dict_on_npu) ... ok
test_state_dict_values_match (__main__.TestScriptModuleStateDict.test_state_dict_values_match) ... ok
test_state_dict_with_destination (__main__.TestScriptModuleStateDict.test_state_dict_with_destination) ... ok
test_state_dict_with_prefix (__main__.TestScriptModuleStateDict.test_state_dict_with_prefix) ... ok
test_to_device (__main__.TestScriptModuleTo.test_to_device) ... ok
test_to_device_and_dtype (__main__.TestScriptModuleTo.test_to_device_and_dtype) ... ok
test_to_dtype (__main__.TestScriptModuleTo.test_to_dtype) ... ok
test_to_dtype_keyword (__main__.TestScriptModuleTo.test_to_dtype_keyword) ... ok
test_to_no_args_returns_self (__main__.TestScriptModuleTo.test_to_no_args_returns_self) ... ok
test_to_npu_and_dtype (__main__.TestScriptModuleTo.test_to_npu_and_dtype) ... ok
test_to_propagates_to_submodules (__main__.TestScriptModuleTo.test_to_propagates_to_submodules) ... ok
test_to_returns_self (__main__.TestScriptModuleTo.test_to_returns_self) ... ok
test_to_string_device (__main__.TestScriptModuleTo.test_to_string_device) ... ok
test_eval_on_npu (__main__.TestScriptModuleTrainEval.test_eval_on_npu) ... ok
test_eval_propagates_to_submodules (__main__.TestScriptModuleTrainEval.test_eval_propagates_to_submodules) ... ok
test_eval_returns_self (__main__.TestScriptModuleTrainEval.test_eval_returns_self) ... ok
test_eval_sets_training_false (__main__.TestScriptModuleTrainEval.test_eval_sets_training_false) ... ok
test_train_default_is_training (__main__.TestScriptModuleTrainEval.test_train_default_is_training) ... ok
test_train_eval_roundtrip (__main__.TestScriptModuleTrainEval.test_train_eval_roundtrip) ... ok
test_train_on_npu (__main__.TestScriptModuleTrainEval.test_train_on_npu) ... ok
test_train_propagates_to_submodules (__main__.TestScriptModuleTrainEval.test_train_propagates_to_submodules) ... ok
test_train_returns_self (__main__.TestScriptModuleTrainEval.test_train_returns_self) ... ok
test_train_set_false (__main__.TestScriptModuleTrainEval.test_train_set_false) ... ok
test_train_set_true_explicit (__main__.TestScriptModuleTrainEval.test_train_set_true_explicit) ... ok
test_type_float32 (__main__.TestScriptModuleType.test_type_float32) ... ok
test_type_float64 (__main__.TestScriptModuleType.test_type_float64) ... ok
test_type_int32_raises (__main__.TestScriptModuleType.test_type_int32_raises) ... ok
test_type_on_npu (__main__.TestScriptModuleType.test_type_on_npu) ... ok
test_requires_grad_unsupported (__main__.TestScriptModuleUnsupported.test_requires_grad_unsupported) ... ok
test_set_extra_state_raises (__main__.TestScriptModuleUnsupported.test_set_extra_state_raises) ... ok
test_to_empty_unsupported (__main__.TestScriptModuleUnsupported.test_to_empty_unsupported) ... ok
test_xpu_unsupported (__main__.TestScriptModuleUnsupported.test_xpu_unsupported) ... ok
test_zero_grad_backward_chain_on_npu (__main__.TestScriptModuleZeroGrad.test_zero_grad_backward_chain_on_npu) ... ok
test_zero_grad_clears_grads (__main__.TestScriptModuleZeroGrad.test_zero_grad_clears_grads) ... ok
test_zero_grad_no_error (__main__.TestScriptModuleZeroGrad.test_zero_grad_no_error) ... ok
test_zero_grad_set_to_none (__main__.TestScriptModuleZeroGrad.test_zero_grad_set_to_none) ... ok
test_zero_grad_set_to_none_false (__main__.TestScriptModuleZeroGrad.test_zero_grad_set_to_none_false) ... ok

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

OK

结果：67 passed，0 failed。

【CheckList】

代码注释完整
assertRaisesRegex 精确匹配异常信息
使用 torch.jit.script() 作为 canonical 创建方式
PR 标题正确使用类型标签：test(jit)

Pull Request已成功合入, 合并人@ascend-robot

（感谢 TensorLake 的贡献）

T TensorLake

6月4日创建了 pull request，commit a9c5745f

T TensorLake

6月4日关联了issue：【Ascend for PyTorch训练营 API一致性任务】补齐测试用例、API功能对齐、补齐文档（9）

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

阶段	任务名	状态	详情
编译构建	Build_X86	✅	>>>
	Build_ARM	✅	>>>
	Build_LibTorch_x86	✅	>>>
	Build_LibTorch_ARM	✅	>>>
	Build_X86_torchair	🛑	>>>
	Build_ARM_torchair	🛑	>>>
	patch_test	🛑	>>>
恶意代码检查	Antipoison	✅	>>>
编码安全与规范检查	CodeCheck	✅	>>>
	check_error	✅	>>>
	CodeCheck_lintrunner	✅	>>>
开源片段检查	SCA	✅	>>>
开发者测试	UT_X86_Part_01	🛑	>>>
	UT_X86_Part_02	🛑	>>>
	UT_ARM_A3_Part_01	🛑	>>>
	UT_ARM_A3_Part_02	🛑	>>>
	UT_ARM_A2_Part_01	✅	>>>
	UT_ARM_A2_Part_02	✅	>>>
	UT_ARM_A2_Part_03	✅	>>>
	UT_inductor_Part_01	🛑	>>>
	UT_inductor_Part_02	🛑	>>>
	UT_inductor_Part_03	🛑	>>>
	UT_inductor_Part_04	🛑	>>>
	UT_DIST_ARM_Part_01	🛑	>>>
	UT_DIST_ARM_Part_02	🛑	>>>
	UT_DIST_ARM_Part_03	🛑	>>>
	UT_DIST_ARM_Part_04	🛑	>>>
	UT_ARM_A2_Select_Part_01	✅	>>>
	UT_ARM_A2_Select_Part_02	✅	>>>
流水线	PR-pipeline_pytorch	✅	>>>