Ascend NPU Bridge for PyTorch

2.9.0 · active · verified Tue Apr 14

torch-npu is a PyTorch extension that serves as an NPU bridge, adapting the Ascend Neural Network Processing Unit (NPU) to the PyTorch framework. It enables developers to leverage the powerful computational capabilities of Huawei Ascend AI Processors for deep learning training and inference within the PyTorch ecosystem. The current version is 2.9.0, with regular updates aligning with PyTorch releases and Ascend software stacks.

Warnings

breaking `torch-npu` and `torch` versions must be strictly aligned. Installing `torch-npu` will often attempt to install a compatible `torch` version, but manual installation requires careful matching. Mismatches can lead to installation failures or runtime errors.
Fix: Always install `torch` and `torch-npu` with the same major.minor.patch version number (e.g., `torch==2.9.0` with `torch-npu==2.9.0`).
gotcha torch-npu requires pre-installation of Huawei's CANN (Heterogeneous Computing Architecture) and HDK (drivers/firmware). These are system-level components and not Python packages. Ensure the CANN environment variables are sourced before running Python scripts.
Fix: Follow the official Ascend documentation to install CANN and HDK. Source the environment script, typically `source /usr/local/Ascend/ascend-toolkit/set_env.sh`, in your shell session or an activation script.
gotcha `torch.npu.set_device()` can only be called once per Python process. Unlike `torch.cuda.set_device()`, it is not possible to switch between NPU devices or set the default device multiple times within a single Python runtime.
Fix: Set the desired NPU device once at the beginning of your script, or use `tensor.to('npu:X')` directly for device placement without relying on a mutable default device. For multi-device scenarios in a single process, ensure tensors are explicitly placed on their target devices and avoid relying on `set_device` after initial setup.
gotcha Ascend NPUs currently do not support the `torch.float64` (double) data type. If a double tensor is created or implicitly used, it will be automatically cast to `torch.float32` (float).
Fix: Design models and data pipelines to primarily use `torch.float32` or `torch.float16` for NPU operations to avoid implicit type conversions and potential precision issues.
gotcha For distributed training or explicit NPU device selection, environment variables like `ASCEND_RT_VISIBLE_DEVICES` or `HCCL_WHITELIST_DISABLE=1` are often required. Incorrect configuration can lead to devices not being utilized or communication errors.
Fix: Refer to distributed training documentation for NPU (e.g., HCCL backend). Set `ASCEND_RT_VISIBLE_DEVICES` to specify visible NPU cards (e.g., `ASCEND_RT_VISIBLE_DEVICES=0,1`). For some scenarios, `export HCCL_WHITELIST_DISABLE=1` might be necessary.

Install

pip install pyyaml setuptools pip install torch==2.9.0 pip install torch-npu==2.9.0 Basic Installation (with x86 CPU PyTorch)
pip install pyyaml setuptools pip install torch==2.9.0+cpu --index-url https://download.pytorch.org/whl/cpu pip install torch-npu==2.9.0 Installation for x86 (if non-NPU PyTorch needed from specific source)
pip install pyyaml setuptools pip install torch==2.9.0 pip install torch-npu==2.9.0 Installation for Aarch64

Imports

torch_npu
```
import torch
import torch_npu
```
While `torch.npu` functions are exposed via `torch` after initialization, explicitly importing `torch_npu` ensures the NPU backend is properly loaded and initialized.
is_available
```
torch.npu.is_available()
```
Used to check if an Ascend NPU device is accessible.
set_device
```
torch.npu.set_device('npu:0')
```
Sets the current NPU device.

Quickstart

This quickstart demonstrates how to check for NPU availability and perform a basic matrix multiplication on an Ascend NPU. It's crucial to first set up the CANN environment variables before running any NPU-accelerated code.

# Ensure CANN environment variables are sourced (e.g., from .bashrc or executed directly)
# source /usr/local/Ascend/ascend-toolkit/set_env.sh

import torch
import torch_npu # Essential for initializing NPU backend

# Check NPU availability
if torch.npu.is_available():
    print(f"NPU is available. Device count: {torch.npu.device_count()}")
    device = torch.device("npu:0")
    # Example tensor operations on NPU
    x = torch.randn(2, 2).to(device)
    y = torch.randn(2, 2).to(device)
    z = x.mm(y)
    print(f"Tensor on NPU:\n{x}")
    print(f"Result of matrix multiplication on NPU:\n{z}")
else:
    print("NPU is not available, using CPU.")
    device = torch.device("cpu")
    x = torch.randn(2, 2).to(device)
    y = torch.randn(2, 2).to(device)
    z = x.mm(y)
    print(f"Tensor on CPU:\n{x}")
    print(f"Result of matrix multiplication on CPU:\n{z}")

view raw JSON →