在提交新问题之前，请确保您已经在社区中搜索过相关问题，并使用了社区中提供的资源/工具后，仍未找到满意的解决方式。

环境信息

环境信息

• 操作系统：乌班图
• 昇腾硬件信息：800I A2
• CANN软件版本：8.5.0
• 安装的对应软件版本：torch torch-npu 2.7.1、2.9.0、2.10.0、2.11.0、2.12.0

使用场景及问题

【涉及 API】

• torch.nn.Module.load_state_dict

1.1 API 功能说明

torch.nn.Module.load_state_dict(state_dict, strict=True, assign=False) 用于将 state_dict 中保存的 parameter 和 persistent buffer 加载到当前 module 及其所有子 module。state_dict 需要是 Mapping 类型，key 与 Module.state_dict() 生成的参数名、buffer 名保持同一命名规则，子模块参数通过 . 分隔的前缀递归匹配。

该接口执行时会先复制输入 state_dict，保留其中的 _metadata 信息，然后从当前 module 开始递归调用各级子模块的 _load_from_state_dict。每个子模块只处理自身 prefix 范围内的 parameter、persistent buffer 和 extra state；加载前会执行 load state dict pre-hook，递归加载完成后会执行 post-hook。

加载 parameter 和 buffer 时，接口会校验 checkpoint 中对应值是否为 Tensor 或 Tensor-like 对象，并检查 shape 是否与当前 module 中的目标对象一致。默认 assign=False 时，接口通过 copy_ 将 checkpoint tensor 的值写入当前 module 已有的 parameter 或 buffer，保留当前 module 中 tensor 对象及其属性。assign=True 时，接口将 state_dict 中的 tensor 赋给 module 中对应的 parameter 或 buffer，保留 state_dict tensor 的属性；Parameter.requires_grad 仍以当前 module 中的设置为准。

strict=True 时，state_dict 的 key 需要与当前 module 的 state_dict() key 完全匹配。缺失 key、未预期 key、shape mismatch、非 Tensor 输入、copy/assign 过程异常等问题会汇总后抛出 RuntimeError。strict=False 时，missing key 和 unexpected key 不作为加载错误抛出，但仍会在返回值中记录。

接口返回 _IncompatibleKeys，其中包含：

• missing_keys：当前 module 期望存在但输入 state_dict 中缺失的 key。
• unexpected_keys：输入 state_dict 中存在但当前 module 不期望的 key。

该接口本身不负责处理 DDP module. 前缀，前缀消费由 torch.nn.modules.utils.consume_prefix_in_state_dict_if_present 等调用方逻辑完成；完成前缀处理后的 state_dict 再进入 load_state_dict 的正常 key 匹配和递归加载流程。

1.2 社区测试完备性说明

PyTorch 官方为该接口提供独立专项测试文件 pytorch/test/nn/test_load_state_dict.py，并在 pytorch/test/test_nn.py、pytorch/test/nn/test_module_hooks.py 中覆盖 extra state 和 load hook 等关联语义。pytorch/test/nn/test_load_state_dict.py 覆盖范围如下。

官方测试已经覆盖 load_state_dict 的核心语义和主要边界行为，测试完备性充分。本次 NPU 适配不新增测试语义，而是在官方测试结构内将关键设备相关路径迁移到 NPU。

1.3 现有 patch 覆盖情况

原 torch-npu upstream patch：

• pytorch-npu/test_upstream/test/nn/test_load_state_dict.py.patch

原 patch 仅在文件顶部导入 torch_npu 和 torch_npu.contrib.transfer_to_npu。但官方 test/nn/test_load_state_dict.py 本身没有 .cuda()、device="cuda" 或 torch.cuda 调用，因此原 patch 不会把关键测试路径迁移到 NPU，也不能证明 NPU module 或 NPU state_dict tensor 的 load_state_dict 行为。

【版本一致性与分支 patch 说明】

已检查 PyTorch v2.7.1、v2.9.0、v2.10.0、v2.11.0、v2.12.0 中 torch.nn.Module.load_state_dict 的实现和相关测试文件。

1. API 实现差异

torch.nn.Module.load_state_dict 位于 torch/nn/modules/module.py。五个版本间的差异如下。

2. 官方测试文件差异

上述版本差异未改变 load_state_dict 的核心语义，也未破坏官方测试覆盖点。