From 4c9522072e0ba1fc28e2e5ba85052d97b13da7fc Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 08:54:41 -0400
Subject: [PATCH] fix: add BitsAndBytesConfig import + remove bf16 GPU strategy

---
 train.py | 35 ++++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/train.py b/train.py
index 560973e..c32c05f 100644
--- a/train.py
+++ b/train.py
@@ -22,6 +22,7 @@ def train(config_path):
     from transformers import (
         AutoModelForCausalLM,
         AutoTokenizer,
+        BitsAndBytesConfig,
         TrainingArguments,
     )
     from peft import (
@@ -60,54 +61,54 @@ def train(config_path):
         print(f"✗ Failed: {e}")
         
         # --------------------------------------------------------------
-        # Strategy 2: BF16 GPU
+        # Strategy 2: BF16 CPU (model too large for single GPU)
         # --------------------------------------------------------------
-        print("\n[2/4] Trying: bf16 GPU...")
+        print("\n[2/4] Trying: bf16 CPU...")
         try:
             model = AutoModelForCausalLM.from_pretrained(
                 config["base_model"],
                 torch_dtype=torch.bfloat16,
-                device_map="auto",
-                trust_remote_code=True,
+                device_map="cpu",
                 low_cpu_mem_usage=True,
+                trust_remote_code=True,
             )
-            print("✓ Success: bf16 GPU")
+            print("✓ Success: bf16 CPU")
         except Exception as e:
-            errors.append(("bf16 GPU", e))
+            errors.append(("bf16 CPU", e))
             print(f"✗ Failed: {e}")
             
             # ----------------------------------------------------------
-            # Strategy 3: BF16 CPU
+            # Strategy 3: FP16 CPU (fallback)
             # ----------------------------------------------------------
-            print("\n[3/4] Trying: bf16 CPU...")
+            print("\n[3/4] Trying: fp16 CPU...")
             try:
                 model = AutoModelForCausalLM.from_pretrained(
                     config["base_model"],
-                    torch_dtype=torch.bfloat16,
+                    torch_dtype=torch.float16,
                     device_map="cpu",
                     trust_remote_code=True,
                     low_cpu_mem_usage=True,
                 )
-                print("✓ Success: bf16 CPU")
+                print("✓ Success: fp16 CPU")
             except Exception as e:
-                errors.append(("bf16 CPU", e))
+                errors.append(("fp16 CPU", e))
                 print(f"✗ Failed: {e}")
                 
                 # ------------------------------------------------------
-                # Strategy 4: FP16 GPU
-                # ------------------------------------------------------
-                print("\n[4/4] Trying: fp16 GPU...")
+                # Strategy 4: 4-bit AS-IS (already quantized)
+                # ----------------------------------------------------------
+                print("\n[4/4] Trying: 4-bit AS-IS...")
                 try:
                     model = AutoModelForCausalLM.from_pretrained(
                         config["base_model"],
                         torch_dtype=torch.float16,
-                        device_map="auto",
+                        device_map="cpu",
                         trust_remote_code=True,
                         low_cpu_mem_usage=True,
                     )
-                    print("✓ Success: fp16 GPU")
+                    print("✓ Success: 4-bit AS-IS")
                 except Exception as e:
-                    errors.append(("fp16 GPU", e))
+                    errors.append(("4-bit AS-IS", e))
                     print(f"✗ Failed: {e}")
                     msg = "\n".join(
                         f"{name}: {err}" for name, err in errors