fix dockerfile and base builder (#1795) [skip-ci]

axolotl-ai-cloud · Jul 30, 2024 · d4f6a6b · d4f6a6b
1 parent d8d1788
commit d4f6a6b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
@@ -14,16 +14,19 @@ jobs:
         include:
           - cuda: "121"
             cuda_version: 12.1.1
+            cudnn_version: 8
             python_version: "3.10"
             pytorch: 2.3.1
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
           - cuda: "121"
             cuda_version: 12.1.1
+            cudnn_version: 8
             python_version: "3.11"
             pytorch: 2.3.1
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
           - cuda: "124"
             cuda_version: 12.4.0
+            cudnn_version: ""
             python_version: "3.11"
             pytorch: 2.4.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
@@ -52,6 +55,7 @@ jobs:
           labels: ${{ steps.metadata.outputs.labels }}
           build-args: |
             CUDA_VERSION=${{ matrix.cuda_version }}
+            CUDNN_VERSION=${{ matrix.cudnn_version }}
             CUDA=${{ matrix.cuda }}
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch }}

diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
@@ -3,7 +3,7 @@ ARG CUDNN_VERSION="8"
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
 
-FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION as base-builder
+FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
 
 ENV PATH="/root/miniconda3/bin:${PATH}"