【合入来源】

如有社区issue，请关联issue链接
请勿携带内部流程信息（需求链接、问题单、内部issue等）

• 需求
• 问题单
• issue/工单
• 重构优化
• 资料更新

【修改方案】

背景说明

通过torch.jit.script()创建的模型是RecursiveScriptModule实例，与测试目的不符，所以本次测试模型均通过继承torch.jit.ScriptModule创建。

torch.jit.ScriptModule 继承自 torch.nn.Module，在 __init__ 执行完毕后通过 init_then_script 将模块编译为 TorchScript，内部持有 _actual_script_module（RecursiveScriptModule 实例）作为 C++ 侧代理。编译过程只迁移参数/子模块/缓冲区的存储，不替换类上的方法，也不影响 hook 相关属性（_backward_hooks、_backward_pre_hooks 等），因此大部分 hook API 在 ScriptModule 子类上行为与普通 Module 一致。

register_full_backward_hook 例外的原因：该方法内部会设置 self._is_full_backward_hook = True，该属性赋值被 ScriptModule.__setattr__ 代理到 C++ 侧后，因类型不匹配（C++ 侧期望 NoneType）导致 RuntimeError。这是pytorch自身的bug。

API 功能说明

register_full_backward_hook(hook, prepend=False) -> RemovableHandle

在模块上注册反向传播后置 hook。hook 签名为 hook(module, grad_input, grad_output) -> tuple[Tensor] or None，在模块梯度计算完成时被调用，可返回新的 grad_input 替代原有值。prepend=True 时 hook 在已有 hook 之前执行。返回 RemovableHandle 用于移除 hook。

register_full_backward_pre_hook(hook, prepend=False) -> RemovableHandle

在模块上注册反向传播前置 hook。hook 签名为 hook(module, grad_output) -> tuple[Tensor] or None，在模块梯度计算之前被调用，可返回新的 grad_output 影响后续梯度计算。prepend=True 时 hook 在已有 hook 之前执行。返回 RemovableHandle 用于移除 hook。

register_load_state_dict_pre_hook(hook) -> RemovableHandle

在 load_state_dict 调用前触发 hook。hook 签名为 hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None，可用于在加载前对 state_dict 进行预处理。返回 RemovableHandle 用于移除 hook。

register_load_state_dict_post_hook(hook) -> RemovableHandle

在 load_state_dict 调用后触发 hook。hook 签名为 hook(module, incompatible_keys) -> None，incompatible_keys 包含 missing_keys 和 unexpected_keys，可原地修改。返回 RemovableHandle 用于移除 hook。

register_state_dict_pre_hook(hook) -> RemovableHandle

在 state_dict 调用前触发 hook。hook 签名为 hook(module, prefix, keep_vars) -> None，可用于在序列化前执行预处理。返回 RemovableHandle 用于移除 hook。

register_state_dict_post_hook(hook) -> RemovableHandle

在 state_dict 调用后触发 hook。hook 签名为 hook(module, state_dict, prefix, local_metadata) -> None，可原地修改 state_dict。返回 RemovableHandle 用于移除 hook。

测试用例说明

register_full_backward_hook

register_full_backward_pre_hook

register_load_state_dict_pre_hook & register_load_state_dict_post_hook

register_state_dict_pre_hook & register_state_dict_post_hook

【资料变更】

不涉及

【接口变更】

不涉及

【功能验证】

测试日志

========================================
Python代码多环境测试开始
测试脚本: /root/torch-2.11/test_script_module_hooks.py
测试时间: Fri Apr 24 06:37:57 AM UTC 2026
========================================

检测到已激活的虚拟环境，先取消激活...
----------------------------------------
测试环境: torch-2.7.1
虚拟环境: /root/torch-2.7.1/.venv/bin/activate
日志文件: /root/torch-2.11/logs/test_2.7.1_20260424_063757.log
----------------------------------------
测试环境: torch-2.7.1
退出码: 0
状态: 成功 ✅
详细日志: /root/torch-2.11/logs/test_2.7.1_20260424_063757.log

完成测试: torch-2.7.1 (退出码: 0)

----------------------------------------
测试环境: torch-2.8
虚拟环境: /root/torch-2.8/.venv/bin/activate
日志文件: /root/torch-2.11/logs/test_2.8_20260424_063757.log
----------------------------------------
测试环境: torch-2.8
退出码: 0
状态: 成功 ✅
详细日志: /root/torch-2.11/logs/test_2.8_20260424_063757.log

完成测试: torch-2.8 (退出码: 0)

----------------------------------------
测试环境: torch-2.9
虚拟环境: /root/torch-2.9/.venv/bin/activate
日志文件: /root/torch-2.11/logs/test_2.9_20260424_063757.log
----------------------------------------
测试环境: torch-2.9
退出码: 0
状态: 成功 ✅
详细日志: /root/torch-2.11/logs/test_2.9_20260424_063757.log

完成测试: torch-2.9 (退出码: 0)

----------------------------------------
测试环境: torch-2.10
虚拟环境: /root/torch-2.10/.venv/bin/activate
日志文件: /root/torch-2.11/logs/test_2.10_20260424_063757.log
----------------------------------------
测试环境: torch-2.10
退出码: 0
状态: 成功 ✅
详细日志: /root/torch-2.11/logs/test_2.10_20260424_063757.log

完成测试: torch-2.10 (退出码: 0)

----------------------------------------
测试环境: torch-2.11
虚拟环境: /root/torch-2.11/.venv/bin/activate
日志文件: /root/torch-2.11/logs/test_2.11_20260424_063757.log
----------------------------------------
测试环境: torch-2.11
退出码: 0
状态: 成功 ✅
详细日志: /root/torch-2.11/logs/test_2.11_20260424_063757.log

完成测试: torch-2.11 (退出码: 0)

========================================
测试完成总结
总测试环境数: 5
完成时间: Fri Apr 24 06:39:05 AM UTC 2026

成功测试数: 5
失败测试数: 0
所有详细日志保存在: /root/torch-2.11/logs
主日志文件: /root/torch-2.11/logs/test_results_20260424_063757.log
========================================

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常