move CUDA_VERSION_BNB arg inside of stage build scope

axolotl-ai-cloud · May 29, 2023 · e43bcc6 · e43bcc6
1 parent 00323f0
commit e43bcc6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
@@ -16,11 +16,11 @@ jobs:
         include:
           - cuda: cu118
             cuda_version: 11.8.0
-            cuda_version_bnb: 118
+            cuda_version_bnb: "118"
             pytorch: 2.0.0
           - cuda: cu117
             cuda_version: 11.7.0
-            cuda_version_bnb: 117
+            cuda_version_bnb: "117"
             pytorch: 1.13.1
     steps:
       - name: Checkout

diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
@@ -1,5 +1,4 @@
 ARG CUDA_VERSION="11.8.0"
-ARG CUDA_VERSION_BNB="118"
 ARG CUDNN_VERSION="8"
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
@@ -62,6 +61,7 @@ RUN git clone https://github.com/microsoft/DeepSpeed.git && \
 FROM base-builder AS bnb-builder
 
 WORKDIR /workspace
+ARG CUDA_VERSION_BNB="118"
 ENV CUDA_VERSION_BNB=$CUDA_VERSION_BNB
 
 RUN git clone https://github.com/TimDettmers/bitsandbytes.git && \
@@ -71,8 +71,6 @@ RUN git clone https://github.com/TimDettmers/bitsandbytes.git && \
 
 FROM base-builder
 
-ENV CUDA_VERSION_BNB=$CUDA_VERSION_BNB
-
 # recompile apex
 RUN python3 -m pip uninstall -y apex
 RUN git clone https://github.com/NVIDIA/apex