用途:把任意 PyTorch 模型迁移到华为昇腾 NPU 并在 NPU 上实际运行验证,覆盖设备检测逻辑、第三方库兼容性和不支持的 API 等问题。触发方式为用户提出“把这个模型迁移到 NPU”“适配到昇腾”“在 NPU 上跑通模型”等请求。
主要能力:按七阶段推进——目标分析(环境、代码、README、迁移难度评估)、快速尝试、方案设计、代码迁移(设备适配、依赖修复、API 替换)、NPU 验证、调试迭代、输出迁移报告。快速尝试阶段推荐在入口文件头部加入 import torch_npu 与 from torch_npu.contrib import transfer_to_npu,该工具通过 monkey-patch 将 torch.cuda.* 重定向到 torch.npu.*,如 torch.cuda.is_available()、tensor.cuda()、empty_cache、set_device、cuda:0、backend=nccl 到 hccl 等;若不能跑通则手动检查 autocast 设备字符串、is_cuda 判断和第三方库 CUDA 调用。深度迁移时按设备替换规则处理,例如 torch.cuda.is_available() 换 torch_npu.npu.is_available()、.cuda() 换 .npu()、nccl 换 hccl,并提供 get_device 模板。验证阶段要求在 NPU 上运行并看到具体数值输出(如 Loss、Output shape、AUC)才算通过,禁止仅凭 import 成功或开始训练判定成功。
使用方式:先通过 find、grep 搜索依赖文件、README、测试与训练入口并分析代码结构;验证命令示例为 ASCEND_VISIBLE_DEVICES=0,1 ASCEND_RT_VISIBLE_DEVICES=0,1 python run.py --config config.yaml,或 pytest tests/test_model.py -v。调试迭代最多 5 次,超出应向用户报告困难。完成后必须生成迁移报告并归档到 references/cases/<模型名小写>.md,报告需含基本信息、环境信息、迁移过程、验证结果、问题与解决等内容。
适用场景与限制:适合单个 PyTorch 模型迁移;若面对的是包含多个模型的框架(如 TorchEasyRec、FuxiCTR),需先告知用户并询问具体迁移哪个模型。约束包括需确保 NPU 驱动与 torch_npu 正确安装、卡号通过环境变量配置、建议不改原代码而在并行目录或 Git 分支上修改。部分细节(如具体支持矩阵、性能对比数据)暂未提供。