From bf07de4b395e7a1d51848422a8e6776edc4b4543 Mon Sep 17 00:00:00 2001
From: ckvermaAI <ckverma@habana.ai>
Date: Fri, 25 Apr 2025 16:29:33 +0300
Subject: [PATCH] bitsandbytes installation

---
 conftest.py                 | 16 ++++++++++++++++
 tests/test_bnb_inference.py | 16 ----------------
 tests/test_bnb_qlora.py     | 19 +------------------
 3 files changed, 17 insertions(+), 34 deletions(-)

diff --git a/conftest.py b/conftest.py
index 1d3329ce3b..b05956913c 100644
--- a/conftest.py
+++ b/conftest.py
@@ -113,6 +113,22 @@ def token(request):
 
 
 def pytest_configure(config):
+    # Bitsandbytes installation for {test_bnb_qlora.py test_bnb_inference.py} tests
+    # This change will be reverted shortly
+    bnb_tests = any("bnb" in name for name in config.known_args_namespace.file_or_dir)
+    if bnb_tests:
+        import subprocess
+        import sys
+
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "git+https://github.com/bitsandbytes-foundation/bitsandbytes.git@multi-backend-refactor",
+            ]
+        )
     name = ""
     try:
         from optimum.habana.utils import get_device_name
diff --git a/tests/test_bnb_inference.py b/tests/test_bnb_inference.py
index 6e76207726..8898e3bd32 100644
--- a/tests/test_bnb_inference.py
+++ b/tests/test_bnb_inference.py
@@ -42,22 +42,6 @@ def get_model(token: str):
 
 @pytest.mark.skipif("gaudi1" == OH_DEVICE_CONTEXT, reason="execution not supported on gaudi1")
 def test_nf4_quantization_inference(token: str, baseline):
-    try:
-        import subprocess
-        import sys
-
-        subprocess.check_call(
-            [
-                sys.executable,
-                "-m",
-                "pip",
-                "install",
-                "git+https://github.com/bitsandbytes-foundation/bitsandbytes.git@multi-backend-refactor",
-            ]
-        )
-    except subprocess.CalledProcessError:
-        pytest.fail("Failed to install bitsandbytes")
-
     os.environ["PT_HPU_LAZY_MODE"] = "0"
     from optimum.habana.transformers import modeling_utils
 
diff --git a/tests/test_bnb_qlora.py b/tests/test_bnb_qlora.py
index 2cd0122304..75b9304954 100644
--- a/tests/test_bnb_qlora.py
+++ b/tests/test_bnb_qlora.py
@@ -14,11 +14,11 @@
 # limitations under the License.
 
 import os
-import subprocess
 
 import pytest
 import torch
 from datasets import load_dataset
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, DataCollatorForLanguageModeling
 
 from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments
@@ -80,23 +80,6 @@ def get_model(token: str):
 
 @pytest.mark.skipif("gaudi1" == OH_DEVICE_CONTEXT, reason="execution not supported on gaudi1")
 def test_nf4_quantization_finetuning(token: str, baseline):
-    try:
-        import sys
-
-        subprocess.check_call(
-            [
-                sys.executable,
-                "-m",
-                "pip",
-                "install",
-                "peft==0.12.0",
-                "git+https://github.com/bitsandbytes-foundation/bitsandbytes.git@multi-backend-refactor",
-            ]
-        )
-        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
-    except subprocess.CalledProcessError:
-        pytest.fail("Failed to install peft==0.12.0 / bitsandbytes")
-
     os.environ["PT_HPU_LAZY_MODE"] = "0"
     from optimum.habana.transformers import modeling_utils