From be7b589c6b3339509c9aaa551f8aa1e609eb7945 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 00:34:34 -0400
Subject: [PATCH] fix: use BitsAndBytes INT4 for PEFT compatibility

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

diff --git a/training/scripts/train.py b/training/scripts/train.py
index e4874d8..d1c1ab2 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -32,15 +32,17 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model - preserve NVFP4 quantization
+    # Load model with BitsAndBytes INT4 (PEFT supports this)
     print(f"Loading model: {config['base_model']}")
+    from transformers import BitsAndBytesConfig
+    
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        torch_dtype=torch.float16,  # Keep NVFP4 quantized
+        quantization_config=BitsAndBytesConfig(load_in_4bit=True),
         device_map="cpu",  # Load to CPU first
         trust_remote_code=True,
     )
-    print("Model loaded to CPU.")
+    print("Model loaded with INT4 quantization.")
     
     # Add LoRA
     lora_config = LoraConfig(