From 1228a1f38be301bb2f6d7c5e8da5701f3009d298 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:11:14 -0400
Subject: [PATCH] fix: load model first then quantize with BnB

---
 training/scripts/train.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/training/scripts/train.py b/training/scripts/train.py
index a4264f1..e6eed2f 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -32,7 +32,7 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model with QLoRA (4-bit quantization)
+    # Load model (skip broken quantization config)
     print(f"Loading model: {config['base_model']}")
     from transformers import BitsAndBytesConfig
     
@@ -43,14 +43,17 @@ def train(config_path):
         bnb_4bit_use_double_quant=True,
     )
     
+    # Load without quantization config, then apply BnB
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        quantization_config=bnb_config,
-        dtype=torch.bfloat16,
-        device_map="auto",  # Load directly to GPU
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
         trust_remote_code=True,
     )
-    print("Model loaded with QLoRA (4-bit).")
+    # Apply 4-bit quantization after loading
+    from peft import prepare_model_for_kbit_training
+    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
+    print("Model loaded and quantized.")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training