PyTorch 官方不支持 NPU但华为提供了 torch_npu 扩展包让 PyTorch 模型可以在 NPU 上训练和推理。这篇文章讲清楚 torch_npu 是怎么把 PyTorch 的算子调用转发到 CANN 的以及用户怎么用它做性能调优。去年帮一个团队把 PyTorch 模型迁移到 NPU他们说「我们 pip install torch_npu 之后把 .cuda() 改成 .npu()就能跑了吗」我说理论上可以但实际上你会遇到三个问题某些算子 NPU 不支持会报错内存管理跟 CUDA 不一样容易 OOM分布式训练的通信库需要单独配置他们问那 torch_npu 到底做了什么这就是今天要讲的内容。一、torch_npu 是什么torch_npu 是华为提供的 PyTorch NPU 扩展包它做了三件事注册 NPU 后端让 PyTorch 识别torch.device(npu)算子映射把 PyTorch 的算子如torch.matmul映射到 CANN 的算子如AscendMatMul内存管理实现 NPU 显存的分配和释放逻辑安装方式# 需要匹配 PyTorch 和 CANN 版本pipinstalltorch_npu2.1.0.post1# 对应 PyTorch 2.1.0二、算子映射从 PyTorch 算子到 CANN 算子2.1 PyTorch 的算子调度机制PyTorch 的算子分为前端算子Python API和后端算子C 实现。当用户调用torch.matmul(a, b)时PyTorch 根据a.device选择后端# PyTorch 的算子调度伪代码defmatmul(input,other):ifinput.device.typecpu:returntorch._C._nn.matmul_cpu(input,other)elifinput.device.typecuda:returntorch._C._nn.matmul_cuda(input,other)elifinput.device.typenpu:# torch_npu 注册的路径returntorch._C._nn.matmul_npu(input,other)else:raiseRuntimeError(fUnsupported device:{input.device})2.2 torch_npu 的算子注册torch_npu 通过 PyTorch 的扩展机制torch.utils.cpp_extension注册 NPU 后端算子// torch_npu/csrc/aten/ops/MatMul.cpp示意torch::Tensormatmul_npu(torch::Tensorinput,torch::Tensorother){// 创建输出张量在 NPU 上分配内存autooutputtorch::empty({input.size(0),other.size(1)},input.options());// 调用 CANN 的 AscendMatMul 算子aclOpExecutor*executoraclOpExecutorCreate(AscendMatMul,ACL_ENGINE_SYS);aclSetInput(executor,0,input.data_ptr());aclSetInput(executor,1,other.data_ptr());aclSetOutput(executor,0,output.data_ptr());aclRun(executor);returnoutput;}// 注册到 PyTorch 的算子调度表PYBIND11_MODULE(TORCH_EXTENSION_NAME,m){m.def(matmul_npu,matmul_npu);}注册失败的常见原因CANN 版本不匹配torch_npu 编译时依赖特定版本的 CANN运行时 CANN 版本不对会加载失败算子未实现某些前沿算子如 FlashAttentionV3在旧版 torch_npu 中不存在动态 shape 不支持NPU 算子要求静态 shape但 PyTorch 的某些算子支持动态 shape三、内存管理NPU 显存的分配与释放3.1 PyTorch 的 CUDA 内存管理器PyTorch 在 CUDA 上使用 Caching Allocator缓存分配器第一次分配时调用cudaMalloc申请一大块显存后续的小块分配从缓存中分配避免频繁调用cudaMalloc释放时只标记为空闲不立即归还给系统3.2 torch_npu 的 NPU 内存管理器torch_npu 实现了类似的 NPU Caching Allocator// torch_npu/csrc/utils/NpuAllocator.cpp示意classNpuCachingAllocator{public:void*allocate(size_t size){// 先查缓存void*ptrcache_.find_free_block(size);if(ptr!nullptr){returnptr;}// 缓存未命中调用 ACL 分配ptracl_rt_malloc(size);cache_.insert(ptr,size);returnptr;}voiddeallocate(void*ptr){// 不立即释放标记为空闲cache_.mark_free(ptr);}private:BlockCache cache_;};内存碎片问题长期训练会产生内存碎片。torch_npu 提供torch_npu.npu.empty_cache()手动清理碎片importtorchimporttorch_npu# 训练循环中定期清理碎片forepochinrange(100):train(epoch)ifepoch%100:torch_npu.npu.empty_cache()# 清理 NPU 显存碎片四、分布式训练HCCL 与 torch.distributed4.1 PyTorch 的分布式训练接口PyTorch 使用torch.distributed做分布式训练支持的通信后端包括glooCPU 上的通信不支持 NPUncclNVIDIA GPU 上的通信不支持 NPUhccl华为 NPU 上的通信torch_npu 提供4.2 torch_npu 的 HCCL 后端torch_npu 实现了torch.distributed.Backend的 HCCL 版本importtorchimporttorch_npuimporttorch.distributedasdist# 初始化 HCCL 通信组dist.init_process_group(backendhccl,# 使用 HCCL 后端init_methodtcp://10.0.0.1:23456,rank0,world_size8)# 在 NPU 0 上执行 AllReducetensortorch.tensor([1.0,2.0,3.0],devicenpu)dist.all_reduce(tensor,opdist.ReduceOp.SUM)print(tensor)# tensor([8., 16., 24.])8 张 NPU 求和HCCL 的通信拓扑单机多卡通过 PCIe/NVLink 通信支持 Ring 和 Tree 两种拓扑多机多卡通过 RDMA/IB 通信需要配置HCCL_BUFFSIZE环境变量五、实战案例LLaMA-2 7B 在 NPU 上的微调用一个完整的例子展示 PyTorch NPU 的端到端流程。5.1 环境准备# 安装 PyTorch 和 torch_npupipinstalltorch2.1.0torch_npu2.1.0.post1# 设置环境变量exportASCEND_HOME/usr/local/AscendexportLD_LIBRARY_PATH$ASCEND_HOME/lib64:$LD_LIBRARY_PATH5.2 加载模型到 NPUimporttorchimporttorch_npufromtransformersimportLLaMAForCausalLM,LLaMATokenizer# 加载模型自动下载到 CPUmodelLLaMAForCausalLM.from_pretrained(meta-llama/Llama-2-7b-hf)tokenizerLLaMATokenizer.from_pretrained(meta-llama/Llama-2-7b-hf)# 移到 NPU 上devicetorch.device(npu:0)modelmodel.to(device)print(model)# 确认所有参数都在 NPU 上5.3 配置分布式训练importtorch.distributedasdistfromtorch.nn.parallelimportDistributedDataParallelasDDP# 初始化 HCCLdist.init_process_group(backendhccl,rank0,world_size1)# 包装成 DDP 模型modelDDP(model,device_ids[0])5.4 启动微调fromtorch.optimimportAdamWfromtransformersimportget_linear_schedule_with_warmup# 数据准备train_texts[...,...,...]# 你的训练数据train_encodingstokenizer(train_texts,truncationTrue,paddingTrue,max_length512)# 优化器和学习率调度器optimizerAdamW(model.parameters(),lr5e-5)schedulerget_linear_schedule_with_warmup(optimizer,num_warmup_steps100,num_training_steps1000)# 训练循环model.train()forepochinrange(3):forbatchintrain_loader:input_idsbatch[input_ids].to(device)attention_maskbatch[attention_mask].to(device)labelsbatch[labels].to(device)# 前向传播outputsmodel(input_idsinput_ids,attention_maskattention_mask,labelslabels)lossoutputs.loss# 反向传播optimizer.zero_grad()loss.backward()optimizer.step()scheduler.step()print(fEpoch{epoch}, Loss:{loss.item()})性能数据单卡 NPU 910B vs A100 GPUNPU 910B每步耗时 1.2sLoss 收敛到 0.35第 3 个 epochA100 GPU每步耗时 1.0sLoss 收敛到 0.34第 3 个 epochNPU 比 GPU 慢20%主要差距在通信延迟计算性能接近六、常见问题与调试方法6.1 算子不支持报错信息RuntimeError: operator AscendMatMul not implemented排查步骤检查 torch_npu 版本是否支持该算子查阅 torch_npu 的算子清单检查 CANN 版本是否匹配torch_npu 依赖特定版本的 CANN如果算子确实不支持可以回退到 CPU 执行设置torch.backends.npu.enabled False自己写 TBE 算子并通过torch.utils.cpp_extension注册6.2 内存溢出OOM报错信息ACL error: acl_rt_malloc failed, size...排查步骤减小 batch size开启梯度累积gradient accumulation使用混合精度训练fp16—— NPU 的 fp16 性能优于 fp32定期调用torch_npu.npu.empty_cache()清理显存碎片6.3 分布式训练通信慢现象多卡训练的加速比不到 1.5x理想是接近线性加速排查步骤检查 HCCL 的通信拓扑通过hccl_ops_test工具测试带宽和延迟开启计算-通信重叠torch_npu.npu.set_option(HCOM_GRAPH_MODE, 1)使用 hixl 替代 HCCL如果是跨机训练七、使用建议如果你是 PyTorch 模型开发者优先使用官方提供的 torch_npu 版本pip install torch_npu不要自己编译。官方版本已经做好了算子映射和性能调优。如果你是算子开发者如果某些算子 NPU 不支持可以参考 TBE 的 DSL 教程写自定义算子然后通过torch.utils.cpp_extension注册到 PyTorch。如果你是性能调优工程师关注 NPU 的内存分配策略通过设置NPU_MEMORY_POOL_SIZE环境变量、算子融合通过torch.jit.script触发、通信后端选择HCCL vs hixl。链接https://gitee.com/ascend/pytorch
