From b300235be04e8c0e684db90e3e828219ae9b98f8 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:53:43 -0400
Subject: [PATCH] fix: load 4-bit model AS-IS without additional quantization

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

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 0c51dab..1c83c52 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -36,15 +36,15 @@ def train(config_path):
     print(f"Loading model: {config['base_model']}")
     from transformers import BitsAndBytesConfig, AutoConfig
     
-    # Skip quantization - use DeepSpeed ZeRO-3 with CPU offload
+    # Load 4-bit model AS-IS (don't apply additional quantization)
     print(f"Loading model: {config['base_model']}")
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        torch_dtype=torch.bfloat16,
-        device_map="cpu",  # Load to CPU first
+        torch_dtype=torch.float16,
+        device_map="auto",
         trust_remote_code=True,
     )
-    print("Model loaded to CPU. DeepSpeed will distribute.")
+    print("Model loaded (4-bit). DeepSpeed will distribute.")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training