已合并
fix: pass map_location to _legacy_load in torch_npu.utils.load for torch.compile cache #43274
wuyouqi1创建于 27 天前
fix: pass map_location to _legacy_load in torch_npu.utils.load for torch.compile cache #43274
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 wuyouqi1 的贡献）
atomgit-bot
27 天前 评论：
27 天前 评论：
变更摘要
此 PR 恢复了之前被回退的改动，修复 torch_npu.utils.load 在 legacy（非 zip）加载分支中对 map_location 参数的透传问题。原先代码在调用 _legacy_load 时将 map_location 硬编码为 "cpu"，导致在 torch.compile 缓存加载等场景下无法将数据正确映射到 NPU 目标设备。本次修改将 _legacy_load 的第二个参数从 "cpu" 改为 map_location，并恢复对应的参数化单测用例，覆盖 None、Callable、Dict 三种 map_location 类型。
主要改动
torch_npu/utils/serialization.py的load函数：在 legacy 加载分支中将_legacy_load(opened_file, "cpu", ...)改为_legacy_load(opened_file, map_location, ...)，使加载结果按调用方传入的map_location正确映射到目标设备test/npu/test_serialization.py新增测试用例test_legacy_load_maplocation：使用@parametrize("map_kind", ["none", "callable", "dict"])参数化驱动，分别覆盖map_location=None、Callable 映射、Dict 映射三种场景，验证 legacy 加载后张量设备正确为npu:0test/npu/test_serialization.py测试框架适配：新增导入parametrize和instantiate_parametrized_tests，并在模块末尾调用instantiate_parametrized_tests(TestSerialization)以激活参数化测试
ascend-robot
27 天前 评论：
27 天前 评论：
atomgit-bot
27 天前 评论：
27 天前 评论：
代码审查
审查总结
变更文件及审查结果：
| 文件 | 审查结论 |
|---|---|
torch_npu/utils/serialization.py |
发现 1 个 P3 问题（警告文本与修复后行为不一致） |
test/npu/test_serialization.py |
无问题 |
按优先级统计：
- P0：0
- P1：0
- P2：0
- P3：1
整体风险评估：低风险。
核心修复（serialization.py line 339："cpu" → map_location）逻辑正确，解决了 torch.compile 缓存在 NPU 上加载时因 map_location=None 被强制改写为 "cpu" 导致 tensor 未能恢复到 NPU 设备的问题。新增的 test_legacy_load_maplocation 参数化测试对 None、callable、dict 三种 map_location 类型均有覆盖，断言正确。
唯一值得关注的是遗留的警告文本（line 315-321）声称 callable/dict 类型不支持且会加载到 CPU，与修复后的实际行为矛盾，建议后续同步更新。
⚠️ 已识别出整体风险，但无法提取行内评论，请参考整体评估。
27 天前 添加了label：ascend-cla/yes
此处折叠了42条消息 查看更多
25 天前 添加了label：approved
ffmh
25 天前 评论：
25 天前 评论：
/lgtm
25 天前 添加了label：lgtm
25 天前 合入了pull request
ascend-robot
25 天前 评论：
25 天前 评论：
流水线 pytorch_gitcode_PR_multiVersion#13484 [ commitID：6aaec4eb ] 已完成
【合入来源】
关联 issue: https://gitcode.com/Ascend/pytorch/issues/3005 (Fixes #3005)
恢复 torch_npu.utils.load 在 legacy（非 zip）加载分支对
map_location的透传，修复 torch.compile 缓存加载在 NPU 上的功能问题。此前该改动被fix: Revert "modifty torch_npu.utils.load for torch.compile cache use"回退，本次恢复原始修改。【修改方案】
torch_npu/utils/serialization.py的load函数中，当走 legacy 加载分支且未命中(torch.device, str)类型分支时，将_legacy_load(opened_file, "cpu", ...)改为_legacy_load(opened_file, map_location, ...)，使加载结果按map_location正确映射到目标设备，保证 torch.compile 缓存场景下的加载行为正确。test/npu/test_serialization.py恢复test_legacy_load_maplocation用例，覆盖 legacy load 的map_location映射。本修改等价于恢复 "modifty torch_npu.utils.load for torch.compile cache use" 的原始改动，仅涉及上述两个文件。
【资料变更】
不涉及
【接口变更】
不涉及
【功能验证】
torch_npu.utils.load在 legacy 分支下的map_location映射；torch.compile 缓存加载。test/npu/test_serialization.py::TestSerialization::test_legacy_load_maplocation。test_legacy_load_maplocationUT 用例看护。【CheckList】