From 1aa9f7d95201f1f49143f25917c93c69a8b344e5 Mon Sep 17 00:00:00 2001
From: Sunny Liu <sunny19981005@outlook.com>
Date: Mon, 21 Apr 2025 16:49:08 -0400
Subject: [PATCH] update multigpu/test_eval, multigpu/test_llama

---
 tests/e2e/multigpu/test_eval.py  | 12 ++++++++----
 tests/e2e/multigpu/test_llama.py | 30 ++++++++++++++++++++++++------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/tests/e2e/multigpu/test_eval.py b/tests/e2e/multigpu/test_eval.py
index 4989b81df..d9ba982f0 100644
--- a/tests/e2e/multigpu/test_eval.py
+++ b/tests/e2e/multigpu/test_eval.py
@@ -30,8 +30,10 @@ class TestMultiGPUEval:
         cfg = DictDefault(
             {
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "load_in_8bit": False,
-                "load_in_4bit": True,
+                "quantization": {
+                    "backend": "bnb",
+                    "bits": 4,
+                },
                 "strict": False,
                 "sequence_len": 2048,
                 "adapter": "qlora",
@@ -99,8 +101,10 @@ class TestMultiGPUEval:
         cfg = DictDefault(
             {
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "load_in_8bit": False,
-                "load_in_4bit": True,
+                "quantization": {
+                    "backend": "bnb",
+                    "bits": 4,
+                },
                 "strict": False,
                 "sequence_len": 2048,
                 "adapter": "qlora",
diff --git a/tests/e2e/multigpu/test_llama.py b/tests/e2e/multigpu/test_llama.py
index 1ff795bd6..d156d6d7b 100644
--- a/tests/e2e/multigpu/test_llama.py
+++ b/tests/e2e/multigpu/test_llama.py
@@ -171,7 +171,10 @@ class TestMultiGPULlama:
                 "sample_packing": False,
                 "eval_sample_packing": False,
                 "pad_to_sequence_len": True,
-                "load_in_8bit": True,
+                "quantization": {
+                    "backend": "bnb",
+                    "bits": 8,
+                },
                 "adapter": "lora",
                 "lora_r": 8,
                 "lora_alpha": 16,
@@ -249,7 +252,10 @@ class TestMultiGPULlama:
                 "sample_packing": False,
                 "eval_sample_packing": False,
                 "pad_to_sequence_len": True,
-                "load_in_4bit": True,
+                "quantization": {
+                    "backend": "bnb",
+                    "bits": 4,
+                },
                 "adapter": "qlora",
                 "lora_r": 8,
                 "lora_alpha": 16,
@@ -548,7 +554,10 @@ class TestMultiGPULlama:
                 "base_model": "axolotl-ai-co/SmolLM2-135M-bnb-nf4-bf16",
                 "adapter": "qlora",
                 "mean_resizing_embeddings": True,
-                "load_in_4bit": True,
+                "quantization": {
+                    "backend": "bnb",
+                    "bits": 4,
+                },
                 "lora_r": 8,
                 "lora_alpha": 16,
                 "lora_dropout": 0.05,
@@ -648,7 +657,10 @@ class TestMultiGPULlama:
                 "lora_alpha": 16,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "load_in_4bit": True,
+                "quantization": {
+                    "backend": "bnb",
+                    "bits": 4,
+                },
             }
         else:
             adapter = {}
@@ -722,7 +734,10 @@ class TestMultiGPULlama:
                 "lora_alpha": 16,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "load_in_4bit": True,
+                "quantization": {
+                    "backend": "bnb",
+                    "bits": 4,
+                },
             }
         else:
             adapter = {}
@@ -796,7 +811,10 @@ class TestMultiGPULlama:
                 "lora_alpha": 16,
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
-                "load_in_4bit": True,
+                "quantization": {
+                    "backend": "bnb",
+                    "bits": 4,
+                },
             }
         else:
             adapter = {}