From b15bc5c975c9707507a8b4ec8fe00b60923899bd Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 01:19:55 -0400
Subject: [PATCH] feat: detect pre-quantized model and skip re-quantization

---
 train.py | 50 +++++++++++++++++++++++++++++++++-----------------
 1 file changed, 33 insertions(+), 17 deletions(-)

diff --git a/train.py b/train.py
index b261970..fc5bfb8 100644
--- a/train.py
+++ b/train.py
@@ -33,24 +33,40 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load bf16 to CPU with BnB 4-bit quantization
-    print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit to CPU...")
-    from transformers import BitsAndBytesConfig
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.float16,
-    )
+    # Check if model is quantized (has model-XXXXX-of-XXXXX files)
+    import glob
+    safetensor_files = glob.glob(f"{config['base_model']}/*.safetensors")
+    is_quantized = any("of-" in f for f in safetensor_files)
     
-    model = AutoModelForCausalLM.from_pretrained(
-        config["base_model"],
-        quantization_config=bnb_config,
-        device_map="cpu",
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-        low_cpu_mem_usage=True,
-    )
-    print("✓ Model loaded to CPU with BnB 4-bit (~17.5GB)")
+    if is_quantized:
+        print(f"\n[INFO] Loading pre-quantized BnB 4-bit model (no re-quantization)...")
+        model = AutoModelForCausalLM.from_pretrained(
+            config["base_model"],
+            device_map="cpu",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("✓ Model loaded to CPU (BnB 4-bit)")
+    else:
+        # Load bf16 with BnB 4-bit quantization
+        print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit to CPU...")
+        from transformers import BitsAndBytesConfig
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            config["base_model"],
+            quantization_config=bnb_config,
+            device_map="cpu",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("✓ Model loaded to CPU with BnB 4-bit (~17.5GB)")
     
     # Move to GPU
     print("   Moving to GPU 0...")