torch npu version

siliconflow · Dec 2, 2024 · f1c3e94 · f1c3e94
1 parent c2b12aa
commit f1c3e94
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 1 deletion.
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -27,6 +27,11 @@ on:
         required: true
         default: "2.1.0"
         description: "Build pytorch version. Example: 2.3.0"
+      torch_npu_version:
+        type: string
+        required: true
+        default: "2.1.0.post3"
+        description: "Build torch-npu version. Example: 2.1.0.post3"
     secrets:
       GH_PRIV_TOKEN:
         required: true
@@ -52,6 +57,11 @@ on:
         required: true
         default: "2.1.0"
         description: "Build pytorch version. Example: 2.3.0"
+      torch_npu_version:
+        type: string
+        required: true
+        default: "2.1.0.post3"
+        description: "Build torch-npu version. Example: 2.1.0.post3"
 
 jobs:
   build:
@@ -76,6 +86,7 @@ jobs:
       - name: Install dependencies
         run: |
           echo "TORCH_VERSION=${{ inputs.torch_version }}" >> ${GITHUB_ENV}
+          echo "TORCH_NPU_VERSION=${{ inputs.torch_npu_version }}" >> ${GITHUB_ENV}
 
           major=$(echo ${{ inputs.python }} | cut -d. -f1)
           minor=$(echo ${{ inputs.python }} | cut -d. -f2)

diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml
@@ -26,6 +26,11 @@ on:
         required: true
         default: "2.1.0"
         description: "Build pytorch version. Example: 2.1.0"
+      torch_npu_version:
+        type: string
+        required: true
+        default: "2.1.0.post3"
+        description: "Build torch-npu version. Example: 2.1.0.post3"
   workflow_dispatch:
     inputs:
       silicondiff_tag:
@@ -38,6 +43,11 @@ on:
         required: true
         default: "2.1.0"
         description: "Build pytorch version. Example: 2.1.0"
+      torch_npu_version:
+        type: string
+        required: true
+        default: "2.1.0.post3"
+        description: "Build torch-npu version. Example: 2.1.0.post3"
 
 permissions:
   contents: read
@@ -65,6 +75,7 @@ jobs:
       python: ${{ matrix.python }}
       silicondiff_tag: ${{ inputs.silicondiff_tag }}
       torch_version: ${{ inputs.torch_version }}
+      torch_npu_version: ${{ inputs.torch_npu_version }}
     secrets:
       GH_PRIV_TOKEN: ${{ secrets.GH_PRIV_TOKEN }}
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -52,6 +52,7 @@ environment-pass = [
   "PIP_NO_BUILD_ISOLATION",
   "PY",
   "TORCH_VERSION",
+  "TORCH_NPU_VERSION",
 ]
 before-build = [
   "bash {project}/downloads/silicondiff-npu-releases/scripts/install_deps.sh",

diff --git a/scripts/install_deps.sh b/scripts/install_deps.sh
@@ -2,4 +2,4 @@ set -Eeuo pipefail
 which python
 which pip
 
-pip install -U setuptools_scm ninja Cython "torch==${TORCH_VERSION}" -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu --no-cache-dir
+pip install -U setuptools_scm ninja Cython "torch==${TORCH_VERSION}" "torch-npu==${TORCH_NPU_VERSION}" -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu --no-cache-dir