move CUDA_VERSION_BNB arg inside of stage build scope

.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

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