Ascend for PyTorch 训练营学习笔记 续一

茶思屋专题：昇腾（Ascend）for PyTorch 训练营(https://www.chaspark.com/#/s/AscendforPyTorch2026?multi=zh) 第一期

PyTorch 侧 与 Torch-NPU 侧 的分层解耦设计：

• PyTorch 侧（上游社区）：负责维护通用多设备能力。

  • Python 层：FSDP、Inductor、DataLoader 等模块利用动态语言特性，自动适配不同设备后端。
  • C++ 层：算子、AMP、Generator、分布式后端等通过函数注册表获取具体设备实现，保持内核精简。

• Torch-NPU 侧（昇腾实现）：基于 CANN 完成 NPU 核心功能的具体实现，通过 register 注册到 PyTorch。

  • NPU 核心实现（③）：Tensor/Storage/Generator 实现、HCCL 通信、AMP、编译后端等，对接 CANN runtime 与 aclnn 算子。
  • NPU 定制化模块（④）：昇腾专属的 Profiler、扩展日志、二级算子下发流水线。
  • 算子层（⑤）：Native OP + Custom OP。
  • 资源管理（⑥）：Event、Stream、Device、Memory。

PyTorch 只管通用接口，昇腾在 out-of-tree 侧基于 CANN 做完整实现，双方通过注册机制松耦合对接。

Torch-X 生态库原生支持 NPU: 昇腾不仅适配了 PyTorch 本体，还向上游生态库贡献了多设备支持代码：

• torchtune：通过 PyTorch 多设备支持机制，让 torchtune 的模型微调能力原生跑在昇腾 NPU 上。
• torchtitan：两个关键 PR 已合入主线：
  • #672（support 3rd-party backend）：支持第三方后端注册。
  • #706（Support 3rd-party distributed backend）：支持第三方分布式后端注册。

架构上 PyTorch 与 NPU 后端解耦对接，生态上 torchtune / torchtitan 等主流库已通过社区合入原生支持昇腾。

总体逻辑架构与数据流：从上层应用到昇腾 NPU 的 完整数据通路，分为 5 条核心链路：

① 计算类 API：torch.matmul/add 等 → PyTorch ATen Dispatcher → torch_npu → CANN aclnn 算子库。

② 分布式 API：allreduce/alltoall 等 → PyTorch c10d → torch_npu → CANN HCCL 通信库。

③ torch.compile：模型入图后生成 FX Graph → Inductor（昇腾后端）→ Triton/MLIR 等 → CANN Runtime 执行。

④ 资源管理 API：Device/Stream/Event/Memory 等 → torch_npu → CANN Runtime。

⑤ 功能算法类 API：模型构造、自动求导、分布式算法等高阶能力，由 ①~④ 组合支撑。

用户写的原生 PyTorch 代码，通过 Dispatcher 分发到 torch_npu，再统一下沉到 CANN 的算子、通信、运行时三层，最终在 NPU 上执行。