TorchNPU

PyTorch Adapter Plugin for Ascend NPU

English · 中文

Installation · Quick Start · Compatibility · Support · Contributing · Docs · Community

Overview

As the core component of the Ascend for PyTorch community, TorchNPU is a deep learning adapter plugin built by Ascend for PyTorch, enabling the PyTorch framework to directly invoke Ascend NPU and providing developers with the powerful compute capabilities of Ascend AI processors.

Ascend is a full-stack AI computing infrastructure for industry applications and services based on Huawei Ascend AI processors and software. For more information about Ascend, see Ascend Community.

Key Features

TorchNPU fully inherits and reuses extensive mature capabilities from upstream PyTorch, and on this foundation provides deep adaptation and optimization for Ascend NPU. The diagram below illustrates the main functional modules of TorchNPU:

Key Modules:

• Basic Compute: Extensive support for PyTorch native APIs and custom APIs, covering mainstream AI scenarios with a consistent experience, enabling rapid model and algorithm implementation.
• Distributed: Supports accelerating distributed training via FSDP2. Core compute APIs support DTensor, along with collective communication primitives such as AllGather, AllReduce, and AllToAll, as well as point-to-point primitives like Send and Recv.
• Graph Mode: Significantly accelerates model training and inference through "dynamic graph capture + static graph optimization + efficient code generation," and supports offloading execution via NPUGraph. Supported in v2.6.0 and above.
• Debug & Tuning: Supports profiling analysis of compute, communication, and memory usage, with real-time communication anomaly monitoring via WatchDog.
• TorchNPU Core: Supports virtual memory management to reduce memory fragmentation, cross-stream memory reuse optimization in distributed scenarios, and integrates operators and device resources into PyTorch via PrivateUse1.

Latest News

• 📢 [2026-06-30] Advance notice on standardized naming conventions for Ascend for PyTorch community components. 🔗 Learn more
• 📢 [2026-04-30] TorchNPU 26.0.0 released, adding support for PyTorch 2.10.0, Python 3.13, P2P communication group dispatch, DTensor strategy extensions, and more. 🔗 Learn more

Installation

From Binary

Take installing TorchNPU 2.10.0.post2 as an example. Run the following commands for installation. For other versions, please refer to the community download page: TorchNPU Download.

Install CANN

Install CANN 9.0.0. For detailed steps, please refer to the CANN Installation Guide.

Install PyTorch

Run the following command to install PyTorch 2.10.0:

pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cpu

Install TorchNPU

Run the following command to install TorchNPU 2.10.0.post2:

pip install torch-npu==2.10.0.post2

From Source

For detailed steps on compiling TorchNPU, please refer to the Source Installation Guide.

Quick Start

Initialize Environment

# Default path, modify according to your actual installation location
source /usr/local/Ascend/ascend-toolkit/set_env.sh

Run Example

import torch
import torch_npu

x = torch.randn(2, 2).npu()
y = torch.randn(2, 2).npu()
z = x.mm(y)
print(z)

Starting from TorchNPU 2.5.1, import torch_npu is no longer mandatory (auto-registration occurs), but explicit import is still recommended to ensure device initialization.

If the following output appears, the execution is successful:

tensor([[-0.0515,  0.3664],
        [-0.1258, -0.5425]], device='npu:0')

For complete model migration and training tutorials, please refer to the Quick Start Guide.

Community

The Ascend for PyTorch community consists of multiple Special Interest Groups (SIGs), each responsible for development, maintenance, and community collaboration in specific technical areas. Below is a list of all current SIGs. Click the corresponding links for detailed descriptions.

Each SIG holds regular meetings, mailing lists, and contribution guides. Click the corresponding SIG links to view detailed contact information, goals, and participation guidelines. Everyone is welcome to contribute to the community. If you have any questions or suggestions, please submit GitHub Issues. We will reply as soon as possible. Thank you for your support.

Security Note

For system security hardening, recommended user configurations, and file permission controls for TorchNPU, please refer to the Security Note.

Disclaimer

This plugin is for debugging and development purposes only. Users are responsible for ensuring the security of input commands and properly controlling access to data generated during use. By using this plugin, you agree to and accept the above statement.

License

For the license of TorchNPU, please refer to the LICENSE file. For the license of TorchNPU documentation, please refer to the LICENSE file.

Acknowledgments

Thanks for every PR from the community! We welcome developers to contribute code to the TorchNPU plugin!