​**作者**​：昇腾实战派

​**知识地图**​：[https://blog.csdn.net/Lumos\_Lovegood/article/details/161601003](https://blog.csdn.net/Lumos_Lovegood/article/details/161601003)

## 一、背景概述

随着深度学习模型的规模与复杂度持续增长，PyTorch 作为业界主流的深度学习框架，其灵活的动态图机制与丰富的生态工具链深受开发者青睐。然而，在将 PyTorch 部署到昇腾 NPU 硬件平台时，我们面临一系列技术挑战：如何在不破坏 PyTorch 社区演进节奏的前提下，实现高效的硬件适配？如何在动态图与静态图之间取得性能与调试体验的平衡？如何在大规模分布式训练场景中，充分发挥 NPU 的通信与计算能力？

本文从框架架构、编译优化、数据流传递、分布式并行及性能优化五个维度，系统梳理昇腾 NPU 适配 PyTorch 的核心技术方案与实践经验，旨在为开发者提供一份清晰的技术参考。

---

## 二、框架架构与接入机制

### 2.1 PyTorch 核心：C++ 与 Aten 算子库

PyTorch 的核心并非 Python，而是 C++。**Aten** 是 PyTorch 的 C++ 算子库，所有诸如 `torch.add`、`torch.mm` 等高层 API 最终都会落到 Aten 中执行。PyTorch 预留了 **C++ 可注册扩展模块** 与 **Aten 可注册接口** 作为扩展点。Torch-NPU 正是通过将定制化的模块和算子注册到这两处，从而在 PyTorch 框架下实现对昇腾 NPU 的便捷调用。

### 2.2 Torch-NPU 的接入方式与组件构成

- **接入方式**：Torch-NPU 采用 **dispatch / out-of-tree** 方式接入 PyTorch 框架。这种方式既不影响 PyTorch 社区的更新节奏，又能快速兼容和支持新版本 torch，同时保留了充分的自由度和定制化优化空间。

- **组件构成**：Ascend for PyTorch 分为两部分：

  - `torch_npu`：主要负责使 NPU 适配 PyTorch 框架，提供设备管理、算子分发等核心能力。

  - `ascend op-plugin`：为 `torch_npu` 提供运行所需的算子适配文件。该组件不感知 PyTorch 版本，通过统一接口对接多个 torch 版本，实现版本解耦，有利于算子优化与维护。

- **关键后端组件**：

  - **NPU Compile Backend**：负责将模型编译为静态图，提升执行效率。

  - **HCCL Process Group**：分布式通信核心，基于华为 HCCL 通信库实现高效的集合通信与点对点通信。

### 2.3 PyTorch 开发者 API 范畴

PyTorch 为开发者提供的 API 覆盖了深度学习全流程：**数据处理、模型构造、图模型加速、分布式并行算法、迭代训练、推理**。Torch-NPU 在此基础上，针对昇腾硬件特性进行了适配与扩展。

---

## 三、动态图与静态图编译优化

### 3.1 PyTorch 2.x 的沙漏型战略

PyTorch 2.x 引入了沙漏型架构设计：

- **顶端**：统一开发者语言，使 Python 成为 AI 开发的通用语言。

- **中间**：以 `torch.compile` 为核心，Torch Dynamo 将 Python 代码转换为 FX Graph，再由 Torch Inductor 生成底层硬件代码。

- **底端**：以往芯片厂商需要为 PyTorch 每个算子手写适配代码，现在只需重点实现 Inductor 的后端（针对 NPU 的代码生成器），Triton / MLIR 就能自动生成成千上万个算子的高效 NPU 内核。Torch-NPU 因此受益，大幅降低了适配成本。

### 3.2 动态图与静态图模式

- **动态图模式（Eager 模式）**：直接使用 Python 动态语言的能力，例如 `y = model(x) -> loss.backward()` —— 每行 Python 代码立即执行，反向传播自动链式求导。其缺点是框架无法进行全局优化。

- **静态图模式（Graph 模式）**：这是 PyTorch / Torch-NPU 提速的关键。当调用 `torch.compile(model)` 时，PyTorch 先使用 Dynamo 将整个模型前向计算过程“捕获”形成计算图（FX Graph），然后将计算图交给 Inductor 生成底层的优化代码（如 Triton / MLIR），最后在 NPU 上执行。

Torch-NPU 新增了 `npu_graph_tree` 图管理模块，该模块能够获取 PyTorch 传递的 FX 图，并将其映射到 NPU 专属的 CANN 编译器中。在 NPU 上，计算图不会直接交给硬件，而是先经过 **Aten IR**（中间表示），然后由 Torch-NPU 的专用编译器（类似 TorchInductor 的 NPU 后端）将图降级成 NPU 特有的 **TBE（张量加速引擎）算子代码**。最终，一个 `for` 循环可能被融合成 NPU 上的一个高性能 Kernel，大幅减少 CPU 与 NPU 之间的交互开销。

### 3.3 `torch.compile` 当前面临的问题

1. **“黑盒化”带来调试困难**：进入图模式后，报错信息可能变成底层的 IR 错误，难以直接对应到某行 Python 代码。

2. **数值差异（精度问题）**：编译器为加速可能改变算子融合顺序或使用不同的数学库，导致编译后的模型输出与原始 Python 模型存在微小偏差。

3. **断图（Graph Breaks）导致性能不稳定**：若代码中包含 Python 控制流（如 `if` 语句动态改变形状）或使用了 `torch.compile` 不支持的算子，图构建会中断，导致部分走图模式、部分回退到 Eager 模式。这是导致性能极慢的最常见原因。

### 3.4 NPU 对 `torch.compile` 的整体解决方案

1. 使用 **Dynamo** 和 **AOTAutograd** 实现图的捕捉和优化，核心流程复用社区逻辑。

2. 在 Inductor 中，框架会经历 `pre_grad_pass` 和 `post_grad_pass`，自动进行模式匹配（如 Conv + BN、Flash Attention、Pad + MatMul 等），将这些组合融合为一个算子，减少 NPU 读写开销和 Kernel 启动开销。

3. Torch-NPU 既支持 PyTorch 社区的标准 **Triton** 语言，也针对昇腾硬件生成专门的 **昇腾 IR** 表示，最终接入 **Bisheng 编译器** 生成底层代码。NPU Inductor 还支持多级精度自动对比和确定性计算。

4. 为解决断图问题，NPU 借助 **NPUGraph_tree**，将切分出的多个子图分别安全地沉降到 NPU 上执行，并通过图间内存复用机制，使子图之间高效传递数据，最大程度挽回断图带来的性能损失。

---

## 四、API 数据流分类与传递路径

不同 API 的数据流传递路径如下：

1. **计算类 API**：Python API → torch Aten dispatch（调度器查询） → torch_npu → CANN aclnn（NPU 高性能算子库）。

2. **分布式 API**：Python API → torch c10d dispatch（拦截 PyTorch 原生通信模块） → torch_npu → CANN HCCL。

3. **编译类 API**：`torch.compile`（捕获前向计算过程） → 计算 / 分布式 API → FX Graph → Inductor (torch_npu) → Triton 等 → CANN runtime。

4. **资源管理 API**：Python API → torch_npu → CANN runtime。常见示例：`torch_npu.npu.set_device(0)`、`torch_npu.npu.synchronize()`、`torch_npu.npu.empty_cache()`。

5. **功能 / 算法类 API**（例如模型构建、分布式算法、自动求导）：这类 API 往往综合了上述 1~4 类 API 的能力。

---

## 五、分布式并行架构与启动优化

### 5.1 PyTorch-NPU 的分布式并行架构

- **加速库（TorchTitan、TorchForge、VeScele）**：本质上不是重写底层逻辑，而是组合多种 PyTorch 原生的分布式能力，形成一套完整的 LLM 训练框架（如配置 FSDP 和张量并行策略）。

- **原生分布式策略层**：包含 FSDP / FSDP2、DDP、TP、PP、Context Parallel。

- **基础并行组件**：

  - Dtensor 的泛化支持

  - 计算 API 能自动完成通信

  - 基于昇腾硬件规格，优化策略传播模型

- **通信基础设施层：ProcessGroup**：华为为底层通信（集合通信与点对点通信）提供了 **HCCL** 和 **SHMEM**（用于共享内存）。

### 5.2 TorchTitan 的组织流程

1. 编写 `ParallelDims` 配置类；

2. TorchTitan 读取它，自动创建 `DeviceMesh`；

3. TorchTitan 自动把模型用 `DTensor` 包装，并应用 FSDP2 策略；

4. 最后，TorchTitan 自动开启 `torch.compile` 做图融合和通信掩盖。

### 5.3 多卡模型启动流程与昇腾优化

多卡模型在开始训练 / 推理前都需经历以下流程：

`torchrun 拉起多进程 → 创建 TCP 通信链（交换 Rank 信息） → 创建 ProcessGroup → 模型初始化 → 执行分布式并行算法 → 通信原语下发 → 开始训练 / 推理`

昇腾针对不同阶段进行了专项优化：

#### （1）初始建链时间

- **问题**：PyTorch 原生建链方式（依赖 TCPStore）在主节点和所有子节点之间需要握手、交换 Rank 信息，建立成千上万个连接，过程串行且同步，耗时巨大。

- **解决方案**：Torch-NPU 采用 **“异步建链、分层建链”** 策略——不再所有节点都直连主节点，而是先按机架、按组分层建链；同时利用 NPU 的 **DMA（直接存储访问）** 能力，让建链过程不阻塞主线程计算，大幅减少大规模集群拉起的时间。

#### （2）内存使用与算法通信

- **问题**：FSDP2 在训练过程中频繁 AllGather，导致通信压力巨大；频繁的内存分配与释放产生显存碎片；量切分策略（Sharding）若不匹配 NPU 的 DMA 特性，通信带宽无法跑满。

- **解决方案**：

  - 通过重计算和重新排布通信顺序，消除多余的 AllGather 算子；

  - 提前开辟一大块固定的通信缓存区，供计算复用，避免重复申请和释放导致的显存碎片。

#### （3）底层卡的调度

- **问题**：每次建立通信域（如从全局通信组中拉出一个子组做 TP 通信）需要跨卡协商，消耗大量控制指令。发指令给 NPU 时，若一个通信算子里包含多个小操作，CPU 会频繁下发多次命令，导致 NPU 空闲时间浪费在等待指令上。

- **解决方案**：

  - 创建基于 Rank 映射表的 **无协商通信域**；

  - 将多条细碎的通信指令 **合并为一条聚合指令**，减少 CPU 与 NPU 的指令交互；

  - 数据在 HCCL 通信时，**不经过系统内存的拷入拷出**，直接在 NPU 显存之间走 DMA 直通，消除带宽冲突，内存带宽利用率拉满。

---

## 六、性能优化专项技术

### 6.1 两级流水线下发算子（CPU 与 NPU 并行）

在普通 Eager 模式中，CPU 需要等待将算子任务完全下发到 NPU 设备后，才能创建新的算子任务。Torch-NPU 在 CPU 和 NPU 之间设置了 **两级流水线**，二级（独立下发进程）负责将 enqueue 指令打包并传递给 NPU 驱动，有效减少 CPU 和 NPU 的等待时间。

### 6.2 绑核与 NUMA 亲和性

- Torch-NPU 会自动完成业务进程绑核，将关键的主线程和下发线程 **强行固定在几个特定的 CPU 核心上**，避免频繁切换带来的性能损耗。

- **NUMA 亲和性**：在多 CPU 插槽服务器上，绑核操作会强制将 NPU 卡的线程运行在离 NPU 物理距离最近的 CPU 核心上，降低内存访问延迟。

### 6.3 NPUGraph 使能图下沉执行

通过提前编译计算图，将逐个 launch 算子转换为 launch 图，从而大幅减少算子 launch 开销，提升整体执行效率。

---

## 七、总结与展望

本文从框架架构、编译优化、数据流传递、分布式并行及性能优化五个维度，系统阐述了昇腾 NPU 适配 PyTorch 的核心技术方案。通过 out-of-tree 接入、动态图与静态图编译优化、两级流水线下发、绑核与 NUMA 亲和性、NPUGraph 图下沉执行等关键技术，Torch-NPU 在保持与 PyTorch 社区同步演进的同时，充分发挥了昇腾硬件的计算与通信能力。

未来，随着 PyTorch 2.x 生态的持续演进，Torch-NPU 将继续深化在 `torch.compile` 后端、分布式并行策略、内存管理等方面的优化，为开发者提供更高效、更易用的昇腾 AI 计算体验。