[docker] update npu dockerfile (#9407)

Co-authored-by: frozenleaves <frozen@Mac.local>
This commit is contained in:
魅影 2025-11-05 18:28:32 +08:00 committed by GitHub
parent eaf963f67f
commit 8edd2622ce
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 7 additions and 3 deletions

View File

@ -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

View File

@ -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"],