From b41d99f956b62c0a9e4cd2acb38ca6f07694c7eb Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:12:57 -0400
Subject: [PATCH] fix: remove broken quantization_config before loading

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

diff --git a/training/scripts/train.py b/training/scripts/train.py
index e6eed2f..70a27a4 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -34,7 +34,11 @@ def train(config_path):
     
     # Load model (skip broken quantization config)
     print(f"Loading model: {config['base_model']}")
-    from transformers import BitsAndBytesConfig
+    from transformers import BitsAndBytesConfig, AutoConfig
+    
+    # Remove broken quantization config
+    model_config = AutoConfig.from_pretrained(config["base_model"])
+    model_config.quantization_config = None
     
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
@@ -43,17 +47,15 @@ def train(config_path):
         bnb_4bit_use_double_quant=True,
     )
     
-    # Load without quantization config, then apply BnB
+    # Load with fresh BnB config
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
+        quantization_config=bnb_config,
         torch_dtype=torch.bfloat16,
         device_map="auto",
         trust_remote_code=True,
     )
-    # 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.")
+    print("Model loaded with QLoRA (4-bit).")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training