From 088521094e45d321d308b83dc73439acc18698a4 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 15:57:42 -0400
Subject: [PATCH] feat: add 4 loading strategies (no BnB for already-quantized
 model)

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

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 0e25155..b30150a 100644
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -32,16 +32,63 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model AS-IS (already 4-bit quantized with CompressedTensors)
+    # Load model - try multiple strategies
     print(f"\n[INFO] Loading {config['base_model']}...")
     
-    model = AutoModelForCausalLM.from_pretrained(
-        config["base_model"],
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-    print("✓ Success: Model loaded successfully")
+    # Strategy 1: 4-bit AS-IS (already quantized)
+    print("\n[1/4] Trying: 4-bit AS-IS...")
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            config["base_model"],
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+        print("✓ Success: 4-bit AS-IS")
+    except Exception as e:
+        print(f"✗ Failed: {e}")
+        
+        # Strategy 2: 4-bit to CPU
+        print("\n[2/4] Trying: 4-bit to CPU...")
+        try:
+            model = AutoModelForCausalLM.from_pretrained(
+                config["base_model"],
+                torch_dtype=torch.float16,
+                device_map="cpu",
+                trust_remote_code=True,
+            )
+            print("✓ Success: 4-bit to CPU")
+        except Exception as e:
+            print(f"✗ Failed: {e}")
+            
+            # Strategy 3: bf16 to CPU
+            print("\n[3/4] Trying: bf16 to CPU...")
+            try:
+                model = AutoModelForCausalLM.from_pretrained(
+                    config["base_model"],
+                    torch_dtype=torch.bfloat16,
+                    device_map="cpu",
+                    low_cpu_mem_usage=True,
+                    trust_remote_code=True,
+                )
+                print("✓ Success: bf16 to CPU")
+            except Exception as e:
+                print(f"✗ Failed: {e}")
+                
+                # Strategy 4: bf16 auto
+                print("\n[4/4] Trying: bf16 auto...")
+                try:
+                    model = AutoModelForCausalLM.from_pretrained(
+                        config["base_model"],
+                        torch_dtype=torch.bfloat16,
+                        device_map="auto",
+                        low_cpu_mem_usage=True,
+                        trust_remote_code=True,
+                    )
+                    print("✓ Success: bf16 auto")
+                except Exception as e:
+                    print(f"✗ Failed: {e}")
+                    raise RuntimeError("All loading strategies failed!")
     
     # Add LoRA
     lora_config = LoraConfig(