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 扩展包，它做了三件事：

1. 注册 NPU 后端：让 PyTorch 识别 torch.device("npu")
2. 算子映射：把 PyTorch 的算子（如 torch.matmul）映射到 CANN 的算子（如 AscendMatMul）
3. 内存管理：实现 NPU 显存的分配和释放逻辑

安装方式：

# 需要匹配 PyTorch 和 CANN 版本
pip install torch_npu==2.1.0.post1  # 对应 PyTorch 2.1.0

二、算子映射：从 PyTorch 算子到 CANN 算子

2.1 PyTorch 的算子调度机制

PyTorch 的算子分为前端算子（Python API）和后端算子（C++ 实现）。当用户调用 torch.matmul(a, b) 时，PyTorch 根据 a.device 选择后端：

# PyTorch 的算子调度（伪代码）
def matmul(input, other):
    if input.device.type == "cpu":
        return torch._C._nn.matmul_cpu(input, other)
    elif input.device.type == "cuda":
        return torch._C._nn.matmul_cuda(input, other)
    elif input.device.type == "npu":  # torch_npu 注册的路径
        return torch._C._nn.matmul_npu(input, other)
    else:
        raise RuntimeError(f"Unsupported device: {input.device}")

2.2 torch_npu 的算子注册

torch_npu 通过 PyTorch 的扩展机制（torch.utils.cpp_extension）注册 NPU 后端算子：

// torch_npu/csrc/aten/ops/MatMul.cpp（示意）
torch::Tensor matmul_npu(torch::Tensor& input, torch::Tensor& other) {
    // 创建输出张量（在 NPU 上分配内存）
    auto output = torch::empty({input.size(0), other.size(1)}, input.options());
    
    // 调用 CANN 的 AscendMatMul 算子
    aclOpExecutor* executor = aclOpExecutorCreate("AscendMatMul", ACL_ENGINE_SYS);
    aclSetInput(executor, 0, input.data_ptr());
    aclSetInput(executor, 1, other.data_ptr());
    aclSetOutput(executor, 0, output.data_ptr());
    aclRun(executor);
    
    return output;
}

// 注册到 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（示意）
class NpuCachingAllocator {
public:
    void* allocate(size_t size) {
        // 先查缓存
        void* ptr = cache_.find_free_block(size);
        if (ptr != nullptr) {
            return ptr;
        }
        
        // 缓存未命中，调用 ACL 分配
        ptr = acl_rt_malloc(size);
        cache_.insert(ptr, size);
        return ptr;
    }
    
    void deallocate(void* ptr) {
        // 不立即释放，标记为空闲
        cache_.mark_free(ptr);
    }
    
private:
    BlockCache cache_;
};

内存碎片问题：长期训练会产生内存碎片。torch_npu 提供 torch_npu.npu.empty_cache() 手动清理碎片：

import torch
import torch_npu

# 训练循环中定期清理碎片
for epoch in range(100):
    train(epoch)
    if epoch % 10 == 0:
        torch_npu.npu.empty_cache()  # 清理 NPU 显存碎片

四、分布式训练：HCCL 与 torch.distributed

4.1 PyTorch 的分布式训练接口

PyTorch 使用 torch.distributed 做分布式训练，支持的通信后端包括：

• gloo：CPU 上的通信（不支持 NPU）
• nccl：NVIDIA GPU 上的通信（不支持 NPU）
• hccl：华为 NPU 上的通信（torch_npu 提供）

4.2 torch_npu 的 HCCL 后端

torch_npu 实现了 torch.distributed.Backend 的 HCCL 版本：

import torch
import torch_npu
import torch.distributed as dist

# 初始化 HCCL 通信组
dist.init_process_group(
    backend="hccl",        # 使用 HCCL 后端
    init_method="tcp://10.0.0.1:23456",
    rank=0,
    world_size=8
)

# 在 NPU 0 上执行 AllReduce
tensor = torch.tensor([1.0, 2.0, 3.0], device="npu")
dist.all_reduce(tensor, op=dist.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_npu
pip install torch==2.1.0 torch_npu==2.1.0.post1

# 设置环境变量
export ASCEND_HOME=/usr/local/Ascend
export LD_LIBRARY_PATH=$ASCEND_HOME/lib64:$LD_LIBRARY_PATH

5.2 加载模型到 NPU

import torch
import torch_npu
from transformers import LLaMAForCausalLM, LLaMATokenizer

# 加载模型（自动下载到 CPU）
model = LLaMAForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = LLaMATokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 移到 NPU 上
device = torch.device("npu:0")
model = model.to(device)

print(model)  # 确认所有参数都在 NPU 上

5.3 配置分布式训练

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化 HCCL
dist.init_process_group(backend="hccl", rank=0, world_size=1)

# 包装成 DDP 模型
model = DDP(model, device_ids=[0])

5.4 启动微调

from torch.optim import AdamW
from transformers import get_linear_schedule_with_warmup

# 数据准备
train_texts = ["...", "...", ...]  # 你的训练数据
train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=512)

# 优化器和学习率调度器
optimizer = AdamW(model.parameters(), lr=5e-5)
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=100, num_training_steps=1000)

# 训练循环
model.train()
for epoch in range(3):
    for batch in train_loader:
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)
        
        # 前向传播
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()
        
        print(f"Epoch {epoch}, Loss: {loss.item()}")

性能数据（单卡 NPU 910B vs A100 GPU）：

• NPU 910B：每步耗时 1.2s，Loss 收敛到 0.35（第 3 个 epoch）
• A100 GPU：每步耗时 1.0s，Loss 收敛到 0.34（第 3 个 epoch）
• NPU 比 GPU 慢 20%（主要差距在通信延迟，计算性能接近）

六、常见问题与调试方法

6.1 算子不支持

报错信息：RuntimeError: operator AscendMatMul not implemented

排查步骤：

1. 检查 torch_npu 版本是否支持该算子（查阅 torch_npu 的算子清单）
2. 检查 CANN 版本是否匹配（torch_npu 依赖特定版本的 CANN）
3. 如果算子确实不支持，可以：
   • 回退到 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

---