## 【合入来源】 Fork: TensorLake/torch-npu 分支: `test/script-module-npu-v2-10-0` → Ascend/pytorch:`v2.10.0` ## 【修改方案】 本 PR 为 [#1861](https://gitcode.com/Ascend/pytorch/issues/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)` | 白名单 | 支持多参数形式，递归传播 | | `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.10.0` + `torch_npu 2.10.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】 - [x] 代码注释完整 - [x] `assertRaisesRegex` 精确匹配异常信息 - [x] 使用 `torch.jit.script()` 作为 canonical 创建方式 - [x] PR 标题正确使用类型标签：`test(jit)`