From 0258dc14d082bd68b6469b677338558ad20297e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B5=AE=E6=A2=A6?=
 <46097299+frozenleaves@users.noreply.github.com>
Date: Tue, 10 Mar 2026 19:37:27 +0800
Subject: [PATCH] [docker] update npu docker (#10268)

Co-authored-by: frozenleaves <frozen@Mac.local>
---
 docker/docker-npu/Dockerfile         | 10 ++++++----
 docker/docker-npu/docker-compose.yml |  2 +-
 requirements/npu.txt                 |  2 +-
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/docker/docker-npu/Dockerfile b/docker/docker-npu/Dockerfile
index 364774f85..034084cb5 100644
--- a/docker/docker-npu/Dockerfile
+++ b/docker/docker-npu/Dockerfile
@@ -1,6 +1,6 @@
 # https://hub.docker.com/r/ascendai/cann/tags
 
-ARG BASE_IMAGE=quay.io/ascend/cann:8.3.rc2-910b-ubuntu22.04-py3.11
+ARG BASE_IMAGE=quay.io/ascend/cann:8.5.1-910b-ubuntu22.04-py3.11
 FROM ${BASE_IMAGE}
 
 # Installation arguments
@@ -33,9 +33,11 @@ RUN pip config set global.index-url "${PIP_INDEX}" && \
 COPY . /app
 
 # Install torch-npu
-RUN pip uninstall -y torch torchvision torchaudio && \
-    pip install --no-cache-dir "torch==2.7.1" "torch-npu==2.7.1" "torchvision==0.22.1" "torchaudio==2.7.1" --index-url "${PYTORCH_INDEX}" && \
-    pip install --no-cache-dir -e . --no-build-isolation && \
+RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh
+RUN pip uninstall -y torch torchvision torchaudio
+RUN pip install --no-cache-dir -r requirements/npu.txt --index-url "${PYTORCH_INDEX}"
+RUN pip install --no-cache-dir -r requirements/deepspeed.txt
+RUN pip install --no-cache-dir -e . --no-build-isolation && \
     pip install --no-cache-dir -r requirements/metrics.txt --no-build-isolation
 
 # Set up volumes
diff --git a/docker/docker-npu/docker-compose.yml b/docker/docker-npu/docker-compose.yml
index 8ef8ddd3b..cda51ca99 100644
--- a/docker/docker-npu/docker-compose.yml
+++ b/docker/docker-npu/docker-compose.yml
@@ -33,7 +33,7 @@ services:
       dockerfile: ./docker/docker-npu/Dockerfile
       context: ../..
       args:
-        BASE_IMAGE: quay.io/ascend/cann:8.3.rc2-a3-ubuntu22.04-py3.11
+        BASE_IMAGE: quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11
         PIP_INDEX: https://pypi.org/simple
     container_name: llamafactory-a3
     image: llamafactory:npu-a3
diff --git a/requirements/npu.txt b/requirements/npu.txt
index 054bd93ea..47fd2eb74 100644
--- a/requirements/npu.txt
+++ b/requirements/npu.txt
@@ -1,4 +1,4 @@
 torch==2.7.1
-torch-npu==2.7.1
+torch-npu==2.7.1.post2
 torchvision==0.22.1
 torchaudio==2.7.1