diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index fea0a9277..30a2e7779 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -109,7 +109,7 @@ jobs:
           platforms: linux/amd64,linux/arm64
           file: ./docker/docker-npu/Dockerfile
           build-args: |
-            BASE_IMAGE=quay.io/ascend/cann:8.3.rc2-a3-ubuntu22.04-py3.11
+            BASE_IMAGE=quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11
           push: ${{ github.event_name != 'pull_request' }}
           tags: |
             docker.io/hiyouga/llamafactory:${{ steps.version.outputs.tag }}-npu-a3
diff --git a/.github/workflows/tests_npu.yml b/.github/workflows/tests_npu.yml
index 448272b47..d4a2f0765 100644
--- a/.github/workflows/tests_npu.yml
+++ b/.github/workflows/tests_npu.yml
@@ -38,7 +38,7 @@ jobs:
       cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
 
     container:
-      image: ascendai/cann:8.3.rc2-910b-ubuntu22.04-py3.11
+      image: ascendai/cann:9.0.0-910b-ubuntu22.04-py3.11
       env:
         HF_ENDPOINT: https://hf-mirror.com
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
@@ -65,8 +65,8 @@ jobs:
       - name: Install dependencies
         run: |
           uv venv
-          uv pip install -r requirements/npu.txt
           uv pip install -e .
+          uv pip install -r requirements/npu.txt
           uv pip install -r requirements/dev.txt
 
       - name: Install node
@@ -89,5 +89,7 @@ jobs:
           make build
 
       - name: Test with pytest
+        shell: bash
         run: |
+          source /usr/local/Ascend/ascend-toolkit/set_env.sh
           make test
diff --git a/docker/docker-npu/Dockerfile b/docker/docker-npu/Dockerfile
index 034084cb5..61da0a60b 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.5.1-910b-ubuntu22.04-py3.11
+ARG BASE_IMAGE=quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11
 FROM ${BASE_IMAGE}
 
 # Installation arguments
diff --git a/docker/docker-npu/docker-compose.yml b/docker/docker-npu/docker-compose.yml
index cda51ca99..75ee904a4 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.5.1-a3-ubuntu22.04-py3.11
+        BASE_IMAGE: quay.io/ascend/cann:9.0.0-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 47fd2eb74..bc07bf885 100644
--- a/requirements/npu.txt
+++ b/requirements/npu.txt
@@ -1,4 +1,5 @@
 torch==2.7.1
-torch-npu==2.7.1.post2
+torch-npu==2.7.1.post4
 torchvision==0.22.1
 torchaudio==2.7.1
+decorator