From 357c9fa7814f29c46e38467a7a72e988c5e92651 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 07:47:42 -0400
Subject: [PATCH] fix: use standard QLoRA pattern with BitsAndBytesConfig

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

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 43bdc6a..0bcbe3b 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -36,17 +36,18 @@ def train(config_path):
     print(f"Loading model: {config['base_model']}")
     from transformers import BitsAndBytesConfig
     
-    quantization_config = BitsAndBytesConfig(
+    bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.bfloat16,
         bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
         bnb_4bit_use_double_quant=True,
     )
     
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        quantization_config=quantization_config,
-        device_map="auto",  # Distribute across GPUs
+        quantization_config=bnb_config,
+        dtype=torch.bfloat16,
+        device_map="auto",
         trust_remote_code=True,
     )
     print("Model loaded with QLoRA (4-bit).")