昇腾 PyTorch 插件（torch_npu）
架构与实现深度解析

Ascend/pytorch/
├── torch_npu/              # 🏠 主 Python 包（核心实现）
│   ├── __init__.py         # 入口文件，初始化所有子模块
│   ├── _afd/               # AFD 模块：Ascend Framework Dispatcher（分布式推理调度）
│   ├── _inductor/           # PyTorch Inductor 后端适配
│   │   └── ascend_npu_ir/   # 基于 MLIR 的 IR 和代码生成
│   ├── contrib/             # 自定义算子/模块（融合注意力、量化等）
│   │   ├── function/        # fused_attention, nms, iou 等
│   │   └── module/          # multihead_attention, deform_conv 等
│   ├── npu/                 # NPU 底层实现（amp、graphs、memory、streams 等）
│   │   ├── aclnn/           # ACLNN 算子 API
│   │   └── npugraph_ex/     # NPUGraph 推理优化
│   └── profiler/             # 性能分析工具
├── torchnpugen/             # 代码生成工具——生成自定义算子注册代码
│   ├── autograd/            # 自动微分代码生成
│   └── templates/           # C++/Python 代码模板
├── third_party/             # 第三方依赖
│   ├── acl/                 # ACL 运行时头文件和 stub
│   ├── hccl/                # HCCL（昇腾集合通信库）
│   ├── dcmi/                # DCMI（设备管理接口）
│   └── torch-mlir/          # MLIR 编译器基础设施
└── examples/                # 示例代码

import torch
import torch_npu.npu
import torch_npu._afd
import torch_npu._inductor
from torch_npu import contrib, profiler

# 为 Tensor/nn.Module 打补丁，添加 .npu() 等方法
from torch_npu.utils import _add_tensor_methods, _apply_module_patch

# 注册 inductor 后端
torch._inductor.config.npu_backend = "mlir"  # MLIR 后端

class ScheduleContextHolder:
    """管理分布式推理的调度上下文"""
    def __init__(self,
                 schedule_mode: int,       # 0=调度FFN, 1=调度Attention
                 session_num: int,          # 会话数量
                 micro_batch_num: int,      # 微批次数量
                 selected_expert_num: int,  # MoE 中选中的专家数量
                 expert_num: int,           # 专家总数
                 attn_to_ffn_token_size: int,
                 ffn_to_attn_token_size: int,
                 ...):

# 方式1：config 配置
import torch_npu._inductor
torch._inductor.config.npu_backend = "mlir"
torch.compile(model)(x)

# 方式2：环境变量
import os
os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
torch.compile(model)(x)

# 方式3：compile options
torch.compile(model, options={"npu_backend": "mlir"})(x)

4.4 昇腾 NPU 的融合实现机制

4.4.1 Inductor-MLIR 融合优化流程

在昇腾的 Inductor-MLIR 框架中，融合发生在 MLIR 的多个层级：

融合 Pass（Pass）是编译器中的一个优化步骤，它会遍历 MLIR 图，识别出可以融合的算子模式，然后用融合后的算子替换原有序列。

4.4.2 融合的数学前提

算子能融合，必须满足数据依赖约束：

# ✅ 可以融合：b 仅仅依赖 a，没有其他分支
a = conv(x)
b = relu(a)

# ✅ 可以融合：残差连接，add 只依赖 conv 和 input，可融合
out = relu(conv(x) + x)

# ❌ 不能融合：c 依赖 b 和 a，必须等两者都完成
a = conv(x)
b = relu(a)
c = add(b, some_other_branch)  # 不能和前面的 relu 融合

编译器需要分析这个数据流图（Data Flow Graph），找到"没有分支逃逸"的算子序列，才能安全融合。

4.4.3 contrib 融合算子的手工注册

除了编译器自动融合，torch_npu.contrib 还提供了手工优化的高性能融合算子：

from torch_npu.contrib.module import multihead_attention

# 直接使用融合版 Multi-Head Attention
model = multihead_attention(embed_dim=512, num_heads=8)

这些手工融合算子经过专门优化，通常比编译器自动融合效果更好，适用于对性能敏感的关键路径。

五HCCL 分布式训练详解

本章节深入剖析 HCCL 在多机多卡训练中的通信机制——这是大模型分布式训练的核心基础设施。

5.1 为什么分布式训练需要集合通信？

先理解一个基本矛盾：

• 单卡算力有限：一块昇腾 NPU 的算力再强，也装不下千亿参数大模型的全部参数
• 模型必须分布在多卡甚至多机上：参数、梯度、优化器状态需要分散存储
• 训练是迭代的：每一步迭代都需要"所有卡上的梯度汇到一起求平均、更新参数，再分发下去"

这就产生了通信需求——多张卡之间必须频繁交换数据（梯度、参数）。

5.2 集合通信的四大基本算子

HCCL 提供了 4 种最基础的集合通信（Collective Communication）操作，理解它们是理解分布式训练的基础：

5.2.1 AllReduce —— 最核心的算子

作用：多卡上的数据分别求和，结果广播给所有卡（每个人都得到一份完整的和）

在大模型训练中的意义：梯度 AllReduce 是同步优化器（如 SGD）的核心步骤。想象成：每个人算出自己部分的梯度，然后所有人把梯度累加，再各自用同样的梯度更新参数——只有这样所有人的参数才能保持同步。

AllReduce 的实现算法：常见的有 Ring-AllReduce（环形算法）和 Tree-AllReduce（树形算法）：

• Ring-AllReduce：N 张卡排成一个环，每个卡只和邻居通信，N-1 轮完成。通信量和卡数无关，适合大带宽高延迟网络。
• Tree-AllReduce：逻辑上是二叉树结构，O(log N) 轮完成，适合低延迟网络。

昇腾 HCCL 根据拓扑自动选择最优算法。

5.2.2 Broadcast —— 一发全收

作用：一张卡（root）上的数据复制到所有其他卡

典型使用场景：参数初始化、主卡（rank 0）加载预训练模型后广播给所有卡。

5.2.3 AllGather —— 人人贡献拼成完整图

作用：每张卡贡献自己的数据片段，最终所有人都拿到完整的拼接结果

典型使用场景：流水线并行（Pipeline Parallelism）中，每个 stage 输出了自己部分的激活值，需要拼成完整的激活序列传给下一个 stage。

5.2.4 ReduceScatter —— 先求和再分散

作用：先把所有卡的数据求和，然后按块分散给不同卡（每张卡拿到总和中属于自己的那片）

典型使用场景：和 AllGather 配对使用，构成完整的数据并行梯度同步流程。

5.3 多机多卡通信拓扑

5.3.1 单机多卡（Same-Host Multi-Card）

5.3.2 多机多卡（Multi-Host）

HCCL 会自动感知拓扑：在初始化时探测网卡、交换机拓扑，选择最优的路由策略。

5.4 分布式训练中的 HCCL 调用流程

5.4.1 初始化

import torch
import torch.distributed as dist
import torch_npu  # 自动加载 HCCL

# 初始化分布式环境
dist.init_process_group(
    backend="nccl",  # 在昇腾上实际是 HCCL，但兼容 NCCL 接口
    init_method="env://",
    world_size=8,    # 总卡数
    rank=0           # 当前卡的编号（0~7）
)

# 设置当前设备
torch.cuda.set_device(rank)  # 在昇腾上等效于 torch.npu.set_device()

5.4.2 数据并行训练（最常见的场景）

数据并行是最简单粗暴的并行策略——每个卡有完整的模型副本，各自处理不同的数据 batch，然后在梯度更新前通过 AllReduce 同步梯度：

for data, target in dataloader:
    # 1. 前向传播（每卡独立）
    output = model(data)
    loss = loss_fn(output, target)

    # 2. 反向传播（每卡独立算出本地梯度）
    loss.backward()  # 此时每卡的梯度只是本地 batch 的，不是全局的！

    # 3. 梯度 AllReduce —— 🎯 HCCL 的核心时刻
    for param in model.parameters():
        dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
        param.grad /= world_size

    # 4. 参数更新（每卡独立，但梯度相同 → 参数同步）
    optimizer.step()

5.4.3 集合通信的异步执行

HCCL 支持异步启动——通信可以和其他计算并行执行（overlap）：

# 异步 AllReduce，不阻塞等待结果
handle = dist.all_reduce(param.grad, op=dist.ReduceOp.SUM, async_op=True)

# 立即开始下一个计算（通信和计算 overlap）
loss2 = model2(input)

handle.wait()  # 等通信完成后再使用结果
param.grad /= world_size
optimizer.step()

这叫通信与计算重叠（Computation-Communication Overlap），是分布式训练性能调优的核心手段。

5.5 HCCL 性能优化技巧

5.5.1 梯度通信与计算重叠

最有效的优化手段之一。在反向传播的同时，启动上一次迭代的梯度通信：

# 伪代码示意
for i, (data, target) in enumerate(dataloader):
    output = model(data)
    loss = loss_fn(output, target)

    # 🎯 关键：反向的同时启动通信
    # 反向传播从最后一层往前算，算完最后一层的梯度时，
    # 最后一层的 AllReduce 可以和更早层的计算并行
    loss.backward()

PyTorch 的 DistributedDataParallel（DDP）自动做了这个优化——它在反向传播时将梯度同步启动，极大减少了通信开销。

5.5.2 通信算子融合（Gradient Bucketing）

如果每次通信单个梯度张量，开销很大。DDP 会把多个小梯度张量打包成一个 bucket，一次通信搞定：

昇腾 HCCL 配合 PyTorch DDP 自动做 bucket 优化。

5.5.3 通信带宽优化

六第三方依赖（third_party/）

6.1 ACL（Ascend CL）——昇腾计算库

ACL（Ascend Caffe Library，后来改叫 Ascend CL）是昇腾的运行时库，提供了：

• 张量操作（aclTensor）
• 算子执行（aclOp）
• 内存管理（aclrt）
• 图编译（ge — Graph Engine）

6.2 HCCL —— 集合通信库

HCCL 是昇腾的多卡通信库，相当于 NVIDIA 的 NCCL，提供：

• AllReduce（多卡梯度求和）
• Broadcast（广播）
• AllGather（收集）
• ReduceScatter（分散求和）

在大模型分布式训练中，HCCL 是多 NPU 卡间数据同步的关键。

6.3 DCMI —— 设备管理

DCMI 是昇腾的设备管理接口，负责 NPU 设备的发现、初始化、状态查询等底层操作。

七关键设计思想

7.1 零入侵接入（Zero-Invasive Integration）

torch_npu 设计的核心原则是：最小化对用户代码的修改。理想情况下，用户只需：

import torch
import torch_npu  # 这一行就够了

device = torch.device("npu")
model = MyModel().to(device)  # 正常写法，torch_npu 自动处理

7.2 猴子补丁（Monkey Patch）

PyTorch 的扩展机制决定了插件需要"修补"一些关键类。torch_npu 通过 _add_tensor_methods()、_apply_module_patch() 等函数为 PyTorch 的：

• torch.Tensor — 添加 .npu() 等方法
• torch.nn.Module — 修改 forward() 以支持 NPU 特定优化
• torch.Storage — 支持 NPU 显存管理

7.3 版本匹配策略

PyTorch 版本演进很快，CANN 版本也在迭代，torch_npu 采用了严格的版本对齐策略：

• 分支命名：{PyTorch版本}-{CANN版本}，如 v2.7.1-7.3.0
• PyTorch 2.1.0 ~ 2.8.0 都有对应的插件版本
• 长期支持（LTS）分支维护 3.5 年，常规分支维护 1 年

八快速入门示例

8.1 最简验证（确认安装成功）

import torch
import torch_npu  # 激活插件

x = torch.randn(2, 2).npu()  # 在 NPU 上创建张量
y = torch.randn(2, 2).npu()
z = x @ y                     # 矩阵乘法
print(z)

8.2 混合精度训练

from torch_npu.npu.amp import GradScaler, autocast

scaler = GradScaler()

for data, target in dataloader:
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

8.3 torch.compile 编译加速（MLIR 后端）

import os
os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'

import torch
import torch_npu

model = MyModel().npu()
model = torch.compile(model, options={"npu_backend": "mlir"})
output = model(input_npu)

8.4 分布式数据并行（DDP）训练

import os
import torch
import torch.distributed as dist
import torch_npu

# 启动方式（2机8卡示例）
# torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr=10.0.0.1 train.py

dist.init_process_group(backend="nccl", init_method="env://")
torch.cuda.set_device(dist.get_rank())

model = MyModel().npu()
model = torch.nn.parallel.DistributedDataParallel(model)

for data, target in dataloader:
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()  # DDP 自动做 AllReduce
    optimizer.step()

九支持的硬件（Atlas 全系列）

9.1 Atlas 训练系列产品

9.2 Atlas 推理系列产品

9.3 Atlas 加速卡（Module / Card 形态）

9.4 芯片代际对照表

十技术总结

十一参考资料

• 昇腾社区官网：https://www.hiascend.com/
• CANN 安装指南：https://www.hiascend.com/cann
• PyTorch 官方：https://pytorch.org/
• GitHub 镜像仓库：https://github.com/Ascend/pytorch