TorchNPU

English | 中文

Quick Reference

TorchNPU is maintained by the Atlas PyTorch community
Where to get help

TorchNPU

The plugin is a deep learning adaptation framework based on Atlas, enabling Atlas NPUs to support the PyTorch framework and providing users of the PyTorch framework with the powerful computing power of Atlas AI processors.

Supported Tags and Dockerfile Usage

Tag Specification

Tags follow the format:

<TorchNPU_version>-<chip>-<os>-<python_version>

Field	value	Description
TorchNPU Version	2.10.0	For details, see the version notes in the readme.
Chip	Specific example values can be found in the CANN mirror	chip model identifier
OS	ubuntu22.04 / openeuler24.03	OS distribution used for the base image
Python Version	py3.11	Major Python version pre-installed in the image

Dockerfile build parameters

See dockerfile:dockerfile

Latest parameters	Description	Required	Reference Source	Value
TORCH_VERSION	Full TorchNPU version number	Yes	TorchNPU repo releases	2.10.0
CHIP_ARCH	chip architecture identifier	Yes	CANN image tag rules	Specific example values can be found in the CANN mirror
OS	Base image operating system	Yes	CANN image tag rules	ubuntu / openeuler
OS_VERSION	Operating system version	Yes	CANN image tag rules	22.04 / 24.03
PY_VERSION	Python version pre-installed in base image	Yes	CANN image tag rules	3.11
CANN_VERSION	CANN toolkit version	Yes	CANN base image repo	9.0.0
ARCH	Host hardware architecture	Yes	Environment hardware	arm / x86
PY_TAG	Python package ABI tag (cp + version number)	Yes	Strictly match PY_VERSION	cp311
TORCH_NPU_RELEASE_VERSION	Official TorchNPU release tag (including PyTorch version)	Yes	TorchNPU repo releases	v26.0.0-pytorch2.10.0
MANYLINUX_VER	PyPI package compatible system version	No	torch official wheel spec	manylinux_2_28
PIP_MIRROR_URL	pip installation source URL (Tsinghua mirror by default)	No	PyPI mirror sources	https://pypi.tuna.tsinghua.edu.cn/simple

Parameter Sources

Image tags, operating systems, and version information:CANN Base Image Repository tag
TORCH_NPU_RELEASE_VERSION parameter reference:TorchNPU Official Release Versions
TORCH_VERSION parameter: https://gitcode.com/Ascend/pytorch/releases/download/v26.0.0-pytorch2.10.0/torch_npu-2.10.0-cp310-cp310-manylinux_2_28_aarch64.whl Extract all content between torch_npu- and -cp310

Quick Start

Build TorchNPU Image

Taking the construction of the 2.10.0-a3-ubuntu22.04-py3.11 image as an example:

docker build \
  --build-arg TORCH_VERSION=2.10.0 \
  --build-arg CHIP_ARCH=a3 \
  --build-arg OS=ubuntu \
  --build-arg OS_VERSION=22.04 \
  --build-arg PY_VERSION=3.11 \
  --build-arg CANN_VERSION=9.0.0 \
  --build-arg ARCH=arm \
  --build-arg PY_TAG=cp311 \
  --build-arg TORCH_NPU_RELEASE_TAG=v26.0.0-pytorch2.10.0 \
  -t image_name:tag \
  -f Dockerfile .

Run TorchNPU Container

docker run \
    --name pta_container \
    --device /dev/davinci1 \
    --device /dev/davinci_manager \
    --device /dev/devmm_svm \
    --device /dev/hisi_hdc \
    -v /usr/local/dcmi:/usr/local/dcmi \
    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
    -v /etc/ascend_install.info:/etc/ascend_install.info \
    -it ascend/pta:tag bash

Secondary Development

# Use TorchNPU image as base image and add user software
FROM quay.io/ascend/torch-npu:2.10.0-910b-ubuntu22.04-py3.11 

RUN apt update -y && \
    apt install gcc ...

...

Supported Hardware

Product Examples	Architecture
Atlas 800T A2, Atlas 900 A2 PoD	ARM64 / x86_64
Atlas 800T A3	ARM64 / x86_64
Atlas 300I Pro, Atlas 300V Pro	ARM64 / x86_64

License

See the license information for TorchNPU included in these images.

Like all container images, pre-installed software packages (Python, system libraries, etc.) may be subject to their own licenses.