前言

两年前如果有人跟我说 PyTorch 模型可以在昇腾 NPU 上跑，而且不用改代码，我肯定不信。那时候要在 NPU 上跑 PyTorch，得自己写算子、自己搭通信、自己调内存，工作量不亚于重新实现一个框架。后来 CANN 社区推出了 pytorch 仓库（在 AtomGit 上也叫 PyTorch 适配仓），事情才真正变得可行。这个仓库做的事情和 tensorflow 仓库类似，但面对的挑战不同——PyTorch 的动态图机制让适配更加复杂，但也更灵活。

pytorch 仓库是昇腾 CANN 生态中负责 PyTorch 框架适配的核心仓库。它让 PyTorch 模型可以几乎无感地运行在昇腾 NPU 上，支持训练和推理全流程，包括分布式训练。这篇文章从仓库定位、核心能力、架构位置、与其他仓库的关系几个维度，把这个仓库讲清楚。

1. 这个仓库是做什么的

pytorch 仓库的定位是 PyTorch 框架适配器，属于 CANN 五层架构中第二层"昇腾计算服务层"的 Framework Adaptor 组件。它和 tensorflow 仓库处于同一层，但适配机制完全不同——这是因为 PyTorch 和 TensorFlow 的底层设计理念不同。

TensorFlow 1.x 是静态图模式，适配相对简单——在图优化阶段拦截计算图，做算子替换就行。PyTorch 是动态图模式（eager mode），每次前向传播都动态构建计算图，无法在编译期做全局优化。所以 pytorch 仓库的适配策略是：在算子级别做拦截，把 PyTorch 的 ATen 算子调用映射到 NPU 算子调用。

具体来说，pytorch 仓库解决四个问题：

1. ATen 算子映射——PyTorch 的底层算子库是 ATen（A Tensor library），pytorch 仓库为 ATen 中的算子提供了 NPU 后端实现

2. TorchScript 编译支持——通过 TorchScript 把动态图转成静态图，然后交给 GE 做全局优化

3. 分布式训练支持——集成 HCCL 通信库，支持 DDP（DistributedDataParallel）分布式训练

4. 自动混合精度——提供 AMP（Automatic Mixed Precision）的 NPU 后端实现

import torch
import torch_npu  # WHY: 这一行就把 PyTorch 的后端注册到了 NPU，之后所有 .to("npu") 的张量都会在 NPU 上分配和计算
# 和 CUDA 的 import torch.cuda 类似，但 NPU 后端是独立的实现，不是 CUDA 的封装

# 检查 NPU 是否可用
print(f"NPU 可用: {torch.npu.is_available()}")
print(f"NPU 数量: {torch.npu.device_count()}")
print(f"当前 NPU: {torch.npu.current_device()}")

# 把张量放到 NPU 上
x = torch.randn(3, 3).npu()  # WHY: .npu() 等价于 .to("npu")，把张量从 CPU 拷贝到 NPU
y = torch.randn(3, 3).npu()
z = torch.matmul(x, y)  # WHY: 这个 matmul 会自动在 NPU 上执行，因为两个输入都在 NPU 上
print(z.cpu())  # WHY: .cpu() 把结果从 NPU 拷回 CPU，用于打印或后续 CPU 操作

WHY 讲解：pytorch 仓库的适配方式比 tensorflow 仓库更优雅。tensorflow 仓库需要用户显式配置 NpuOptimizer，而 pytorch 仓库只需要 import torch_npu，之后所有 NPU 相关的操作都通过 torch.npu API 完成，和 torch.cuda 的用法几乎一样。这意味着如果你之前在 GPU 上用 PyTorch，迁移到 NPU 基本就是把 cuda 替换成 npu。这种设计降低了学习成本，也让大量现有的 PyTorch 代码可以几乎零修改地跑在 NPU 上。

2. 核心能力拆解

2.1 ATen 算子映射机制

PyTorch 的 ATen 库定义了所有张量操作的接口，每种操作有多种后端实现：CPU、CUDA、MPS 等。pytorch 仓库添加了 NPU 后端实现。

ATen 算子的分发机制是 dispatch key——每个算子调用时，PyTorch 根据输入张量的 dispatch key（比如 CPU、CUDA、NPU）来选择对应的后端实现。pytorch 仓库注册了 “NPU” 这个 dispatch key，当输入张量在 NPU 上时，算子调用会被路由到 NPU 实现。

映射的覆盖率是关键指标。pytorch 仓库覆盖了 ATen 中 95% 以上的常用算子，包括：

1. 数学运算——add, sub, mul, div, matmul, bmm 等

2. 神经网络算子——conv2d, conv3d, batch_norm, layer_norm, softmax, gelu 等

3. 张量操作——reshape, permute, slice, index_select, gather, scatter 等

4. 随机数——rand, randn, uniform, dropout 等

5. 归约运算——sum, mean, max, min, argmax, argmin 等

import torch
import torch_npu

# 算子映射的优先级
x = torch.randn(1024, 1024).npu()
w = torch.randn(1024, 1024).npu()

# 这个 matmul 会被分发到 NPU 后端
result = torch.matmul(x, w)  # WHY: NPU 后端的 matmul 使用 Cube 单元执行，FP16 下吞吐量约 128 TOPS

# 检查某个算子是否支持 NPU 后端
print(torch.npu.is_support("matmul"))  # WHY: 返回 True 表示 NPU 后端有 matmul 的实现
print(torch.npu.is_support("some_custom_op"))  # WHY: 返回 False 表示需要自定义算子或回退到 CPU

# 查看算子是否回退到了 CPU
torch.npu.set_option("enable_debug", True)  # WHY: 开启调试模式，算子回退到 CPU 时会打印警告
# 在生产环境中关闭，因为调试模式有额外开销

WHY 讲解：is_support 是排查算子兼容性的快速方法。如果返回 False，说明这个算子没有 NPU 后端实现，会被自动回退到 CPU 执行。回退到 CPU 的算子会导致 NPU-CPU-NPU 的数据搬运，严重影响性能。遇到不支持的算子，处理方式有三种：等价替换（用支持的算子组合实现相同功能）、自定义算子（用 Ascend C 编写）、或者接受回退（如果算子不是热点）。调试模式在生产环境中必须关闭，因为每次算子分发都会做额外的检查，开销约 5-10%。

2.2 TorchScript 编译优化

PyTorch 的动态图模式虽然灵活，但无法在编译期做全局优化。pytorch 仓库通过 TorchScript 把动态图转成静态图，然后交给 GE 做全局优化。

TorchScript 是 PyTorch 官方的图捕获工具，它可以把 Python 函数转成静态计算图。pytorch 仓库扩展了 TorchScript 的编译后端，增加了 GE 编译选项。

import torch
import torch_npu

class BERTModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(1024, 1024)
        self.norm = torch.nn.LayerNorm(1024)

    def forward(self, x):
        return self.norm(self.linear(x))

model = BERTModel().npu()

# 方法 1：直接 eager 执行（无全局优化）
output_eager = model(torch.randn(1, 1024).npu())

# 方法 2：TorchScript 编译 + GE 优化
scripted_model = torch.jit.script(model)  # WHY: script 模式会把 Python 代码转成静态图，
# GE 可以在这个静态图上做算子融合、内存规划等全局优化
# trace 模式（torch.jit.trace）也可以，但对控制流的支持不如 script

# 使用 GE 后端编译
with torch.npu.stream(torch.npu.Stream()):
    output_scripted = scripted_model(torch.randn(1, 1024).npu())

WHY 讲解：TorchScript 编译的收益在 BERT-Large 这种大模型上非常明显。eager 模式下，每个算子独立执行，算子间的数据需要写回 HBM 再读出来。TorchScript + GE 编译后，算子可以融合执行，中间结果留在 L1/L2 缓存中。在 BERT-Large 的测试中，TorchScript 编译后的推理延迟比 eager 模式低 30-40%。但 TorchScript 也有局限——它不支持 Python 的一些动态特性（比如动态 shape、条件分支依赖运行时值），所以不是所有模型都能用 TorchScript 编译。

2.3 分布式训练支持

pytorch 仓库的分布式训练基于 HCCL 通信库，支持 PyTorch 的 DDP 模式。配置方式和 CUDA DDP 几乎一样，只需要把 init_method 和 backend 换成 NPU 对应的配置。

import torch
import torch_npu
import torch.distributed as dist

def setup_distributed():
    # 初始化进程组
    dist.init_process_group(
        backend="hccl",  # WHY: 使用 HCCL 后端，不用 nccl。HCCL 是昇腾 NPU 的集合通信库，
        # 接口和 NCCL 类似但底层实现完全不同，针对昇腾达芬奇架构做了优化
        init_method="tcp://10.0.0.1:29500",
        rank=int(os.environ["RANK"]),
        world_size=int(os.environ["WORLD_SIZE"])
    )
    torch.npu.set_device(int(os.environ["LOCAL_RANK"]))  # WHY: 每个进程绑定一张 NPU 卡，
    # LOCAL_RANK 是当前节点内的卡编号，RANK 是全局进程编号

def train():
    model = MyModel().npu()
    model = torch.nn.parallel.DistributedDataParallel(
        model,
        device_ids=[int(os.environ["LOCAL_RANK"])],  # WHY: 指定 DDP 使用的设备 ID，
        # DDP 会在每个设备上复制模型，梯度同步由 HCCL 的 AllReduce 完成
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for data, target in dataloader:
        data = data.npu()
        target = target.npu()
        output = model(data)
        loss = torch.nn.functional.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

WHY 讲解：DDP 在 NPU 上的工作方式和 GPU 上一样——每个进程持有一份模型副本，前向传播独立计算，反向传播时通过 AllReduce 同步梯度。区别在于通信后端从 NCCL 换成了 HCCL。HCCL 的 AllReduce 在昇腾 NPU 上的实现利用了 NPU 之间的高速互联通道（HCCS），带宽和延迟指标与 NVLink 相当。但需要注意，HCCL 的拓扑感知和 NVLink 不同——昇腾服务器的 NPU 互联拓扑是确定性的，不需要像 NCCL 那样做拓扑发现，但也不支持 NVLink 的自适应路由。

3. 在 CANN 架构中的位置

pytorch 仓库在 CANN 五层架构中的位置：

• 上层：PyTorch 用户代码（Python），你写的模型和训练逻辑
• 本层：pytorch 仓库（Framework Adaptor），注册 NPU 后端、提供算子映射
• 下层：GE（图编译）、AOL 算子库（算子实现）、HCCL（通信）、Runtime（执行）

和 tensorflow 仓库的关键区别：tensorflow 仓库在 Grappler 层面拦截计算图，而 pytorch 仓库在 ATen dispatch 层面拦截算子调用。这意味着：

1. tensorflow 仓库可以在图优化阶段做全局优化（算子融合、内存规划）

2. pytorch 仓库在 eager 模式下只能做算子级别的优化，需要 TorchScript 才能做全局优化

3. tensorflow 仓库的适配更"侵入"（需要配置 NpuOptimizer），pytorch 仓库的适配更"透明"（只需 import torch_npu）

和 pytorch 仓库关系最密切的几个仓库：

1. ge——图引擎，TorchScript 编译后的图由 GE 处理

2. hccl——集合通信库，DDP 的通信层

3. ops-nn / ops-math——算子库，pytorch 仓库映射的算子最终由这些仓库实现

4. ascend-transformer-boost (ATB)——Transformer 加速库，提供融合算子的上层封装

5. torchtitan-npu——基于 PyTorch Titan 的大模型训练框架，直接使用 pytorch 仓库

6. cann-recipes-train——训练配方仓库，提供完整的训练示例

4. 效率对比

指标	迁移前（手动适配）	迁移后（pytorch 仓库）	提升
PyTorch 模型迁移时间	3-5 天	0.5-1 天	3-6 倍
训练性能	GPU 1.2 天/epoch	NPU 0.4 天/epoch	3 倍
算子映射覆盖率	手动逐个实现	95%+ 自动映射	大幅降低开发量
DDP 通信延迟	NCCL 120 us	HCCL 95 us	1.3 倍
AMP 混合精度	手动管理 loss scale	自动 AMP	简化开发流程

迁移时间从 3-5 天降到 0.5-1 天，主要因为 pytorch 仓库的 import-and-go 设计——import torch_npu 后，大部分模型只需要把 .cuda() 换成 .npu() 就能跑。训练性能的 3 倍提升来自 Cube 单元对矩阵运算的并行度优势和算子融合。

5. 常见问题与解决思路

5.1 算子不支持

如果遇到不支持的算子，pytorch 仓库会自动回退到 CPU 执行，并打印警告。可以通过 torch.npu.is_support() 检查特定算子是否支持。

解决思路和 tensorflow 仓库一样：等价替换、自定义算子、或接受回退。

5.2 TorchScript 编译失败

TorchScript 对 Python 动态特性支持有限，常见失败原因：

1. 动态 shape 依赖运行时值——比如 x.view(-1, hidden_dim) 中的 -1 依赖运行时 shape

2. 条件分支依赖运行时值——比如 if x.size(0) > 0:

3. 使用了 Python 标准库函数——比如 len(x) 在某些情况下不能被 TorchScript 解析

解决思路：尽量在模型代码中避免这些动态特性，或者用 torch.jit.trace 代替 torch.jit.script（trace 只记录一次前向传播的执行路径，不支持条件分支，但对动态 shape 更友好）。

5.3 显存不足

NPU 的显存管理和 GPU 不同。PyTorch 在 GPU 上使用 caching allocator，NPU 上也提供了类似的机制，但行为有差异。

import torch
import torch_npu

# NPU 显存管理
torch.npu.set_per_process_memory_fraction(0.8, 0)  # WHY: 限制当前进程最多使用 80% 的 NPU 显存，
# 留 20% 给系统和其他进程。如果不设置，单个进程可能占满所有显存，导致其他进程无法分配显存。

# 清空缓存
torch.npu.empty_cache()  # WHY: 释放 PyTorch 缓存池中的空闲显存，但不影响正在使用的张量。
# 和 torch.cuda.empty_cache() 用法一样，但 NPU 的缓存策略略有不同——NPU 会更积极地缓存小张量

WHY 讲解：显存管理在多进程共享 NPU 的场景下特别重要。PyTorch 的 caching allocator 会预分配一大块显存，然后在这个块内做分配。如果两个进程都预分配了大部分显存，就会出现一个进程闲置显存但另一个进程分配不了的情况。set_per_process_memory_fraction 可以限制每个进程的显存用量，避免这种问题。

自然收尾

pytorch 仓库解决的核心问题和 tensorflow 仓库一样——让模型在昇腾 NPU 上跑起来，而且跑得快。但实现路径不同：tensorflow 仓库走的是图优化路线，pytorch 仓库走的是算子分发路线。这意味着 pytorch 仓库的适配更轻量、更透明，但在全局优化上需要 TorchScript 的配合。如果你的 PyTorch 模型结构规整、不需要太多动态特性，TorchScript + GE 编译能给你带来显著的性能提升；如果模型很动态，eager 模式下也能正常跑，只是少了一些全局优化的机会。

---

仓库链接：https://atomgit.com/cann/tensorflow