From 8edd2622cedec2ee85da474a6c3c8ec90b1e329e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=AD=85=E5=BD=B1?= <46097299+frozenleaves@users.noreply.github.com> Date: Wed, 5 Nov 2025 18:28:32 +0800 Subject: [PATCH] [docker] update npu dockerfile (#9407) Co-authored-by: frozenleaves --- docker/docker-npu/Dockerfile | 8 ++++++-- setup.py | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/docker/docker-npu/Dockerfile b/docker/docker-npu/Dockerfile index 7cd5e8f4..7539008b 100644 --- a/docker/docker-npu/Dockerfile +++ b/docker/docker-npu/Dockerfile @@ -1,5 +1,9 @@ # https://hub.docker.com/r/ascendai/cann/tags -ARG BASE_IMAGE=ascendai/cann:8.1.rc1-910b-ubuntu22.04-py3.11 + +# default base image build for A2, if build for A3, using this image: +# ARG BASE_IMAGE=ascendai/cann:8.3.rc1-a3-ubuntu22.04-py3.11 + +ARG BASE_IMAGE=ascendai/cann:8.3.rc1-910b-ubuntu22.04-py3.11 FROM ${BASE_IMAGE} # Installation arguments @@ -31,7 +35,7 @@ RUN pip config set global.index-url "${PIP_INDEX}" && \ # Install torch-npu RUN pip uninstall -y torch torchvision torchaudio && \ - pip install --no-cache-dir "torch-npu==2.5.1" "torchvision==0.20.1" --index-url "${PYTORCH_INDEX}" + pip install --no-cache-dir "torch==2.7.1" "torch-npu==2.7.1" "torchvision==0.22.1" --index-url "${PYTORCH_INDEX}" # Install the requirements COPY requirements.txt /app diff --git a/setup.py b/setup.py index dc79dddf..d7b5ca63 100644 --- a/setup.py +++ b/setup.py @@ -43,7 +43,7 @@ def get_console_scripts() -> list[str]: extra_require = { "torch": ["torch>=2.0.0", "torchvision>=0.15.0"], - "torch-npu": ["torch-npu==2.5.1", "torchvision==0.20.1", "decorator"], + "torch-npu": ["torch==2.7.1", "torch-npu==2.7.1", "torchvision==0.22.1", "decorator"], "metrics": ["nltk", "jieba", "rouge-chinese"], "deepspeed": ["deepspeed>=0.10.0,<=0.16.9"], "liger-kernel": ["liger-kernel>=0.5.5"],