From e43bcc6c4f26c5084b76f10edbe4864e7e5601f0 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Mon, 29 May 2023 13:30:15 -0400
Subject: [PATCH] move CUDA_VERSION_BNB arg inside of stage build scope

---
 .github/workflows/base.yml | 4 ++--
 docker/Dockerfile-base     | 4 +---
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
index 5a3f90992..7a3e2950d 100644
--- 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
index 510d038ec..a61f6d42d 100644
--- 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