torch-npu

library2.10.0pypypiunverified

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.

ai-ml

pip install pyyaml setuptools
pip install torch==2.9.0
pip install torch-npu==2.9.0

package ↗homepage ↗

INSTALL

IMPORT

SIG · TORCH-NPU

torch-npu

ai-mlpythonv2.10.0

Install

—

Import

—

Disk

—

Pass rate

0/ 10

Env Coverage0 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v2.9.0 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

glibc

py 3.10

✕ no_wheel

py 3.11

✕ no_wheel

4/12 runs

py 3.12

✕ no_wheel

4/12 runs

py 3.13

✕ no_wheel

✕ dependency_conflict

py 3.9

✕ no_wheel

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

torch_npu

✓ import torch_npu

✗ import torch.npu # without prior import torch_npu

import torch_npu after torch — this registers the NPU backend. torch.npu functions are then available via torch.

torch

✓ import torch

Always import torch before torch_npu

is_available

✓ torch.npu.is_available()

Check NPU availability before use

set_device

✓ torch.npu.set_device('npu:0')

Can only be called once per process

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}")

Debug

Known issues

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`).

affects: All versions

gotchatorch-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.

affects: All versions

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.

affects: Before 2.1.0, and partially fixed in 2.1.0 and above, but still has limitations compared to CUDA.

gotchaAscend 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.

affects: All versions

gotchaFor 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.

affects: All versions

Upgrade

Version history

2.10.0latest on PyPI · released May 6, 2026

Audit

Dependencies

torchrequiredCore PyTorch library; must be version-aligned with torch-npu.

CANNrequiredHuawei Ascend Heterogeneous Computing Architecture; a system-level prerequisite for NPU operation.

HDKrequiredAscend drivers and firmware; a system-level prerequisite for NPU operation.

pyyamlrequiredRuntime dependency for torch-npu.

setuptoolsrequiredRuntime dependency for torch-npu.

Agent activity

153 hits · last 30 days

node

144

OpenAI (training)

Resources

packagehttps://pypi.org/project/torch-npu/ ↗

homepagehttps://gitcode.com/ascend/pytorch ↗