From 21f17cca691d5df41863184a188b259feada99bb Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Mon, 29 May 2023 00:06:35 -0400
Subject: [PATCH] bnb fixes

---
 .github/workflows/base.yml |  2 +-
 docker/Dockerfile-base     |  2 +-
 scripts/finetune.py        | 18 +++++++++---------
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml
index 571faf771..5a3f90992 100644
--- a/.github/workflows/base.yml
+++ b/.github/workflows/base.yml
@@ -1,4 +1,4 @@
-name: ci-cd
+name: ci-cd-base
 
 on:
   push:
diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
index 4f4431dbe..6399d60ee 100644
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -90,7 +90,7 @@ COPY --from=flash-attn-builder /workspace/flash-attention/csrc/rotary/dist/rotar
 COPY --from=flash-attn-builder /workspace/flash-attention/csrc/layer_norm/dist/dropout_layer_norm-*.whl wheels
 
 RUN pip3 install wheels/deepspeed-*.whl wheels/flash_attn-*.whl wheels/fused_dense_lib-*.whl wheels/xentropy_cuda_lib-*.whl wheels/rotary_emb-*.whl wheels/dropout_layer_norm-*.whl
-RUN cd /workspace/builds/bitsandbytes && python3 setup.py install
+RUN cd /workspace/builds/bitsandbytes && cp bitsandbytes/libbitsandbytes_cuda.so bitsandbytes/libbitsandbytes_cuda${CUDA_VERSION_BNB}.so && python3 setup.py install
 RUN git lfs install --skip-repo
 RUN pip3 install "peft @ git+https://github.com/huggingface/peft.git@main" \
         "accelerate @ git+https://github.com/huggingface/accelerate.git@main" \
diff --git a/scripts/finetune.py b/scripts/finetune.py
index 1d1eb9f95..58f1c0957 100644
--- a/scripts/finetune.py
+++ b/scripts/finetune.py
@@ -178,6 +178,15 @@ def train(
             tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH
         )
 
+    if cfg.debug or "debug" in kwargs:
+        logging.info("check_dataset_labels...")
+        check_dataset_labels(
+            train_dataset.select(
+                [random.randrange(0, len(train_dataset) - 1) for i in range(5)]
+            ),
+            tokenizer,
+        )
+
     if prepare_ds_only:
         logging.info("Finished preparing dataset. Exiting...")
         return
@@ -213,15 +222,6 @@ def train(
         model.save_pretrained(cfg.output_dir)
         return
 
-    if cfg.debug:
-        logging.info("check_dataset_labels...")
-        check_dataset_labels(
-            train_dataset.select(
-                [random.randrange(0, len(train_dataset) - 1) for i in range(5)]
-            ),
-            tokenizer,
-        )
-
     trainer = setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer)
 
     model.config.use_cache = False