From 9d8c2b6cf0829b1293f23958ebace77435135fe9 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 21:38:38 -0400
Subject: [PATCH] refactor: cleaner loading strategies with error tracking

---
 train.py | 57 ++++++++++++++++++++++++++++++++++++++------------------
 1 file changed, 39 insertions(+), 18 deletions(-)

diff --git a/train.py b/train.py
index 0829e88..46744cb 100644
--- a/train.py
+++ b/train.py
@@ -32,68 +32,89 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model - try multiple strategies
+    # Load model with multiple strategies
     print(f"\n[INFO] Loading {config['base_model']}...")
+    errors = []
     
-    # Strategy 1: bf16 model with BnB 4-bit quantization
-    print("\n[1/4] Trying: bf16 with BnB 4-bit...")
+    # ------------------------------------------------------------------
+    # Strategy 1: QLoRA (preferred)
+    # ------------------------------------------------------------------
+    print("\n[1/4] Trying: 4-bit NF4 on GPU...")
     try:
-        from transformers import BitsAndBytesConfig
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             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=bnb_config,
             device_map="auto",
             trust_remote_code=True,
+            low_cpu_mem_usage=True,
         )
-        print("✓ Success: bf16 with BnB 4-bit")
+        print("✓ Success: QLoRA 4-bit")
     except Exception as e:
+        errors.append(("QLoRA 4-bit", e))
         print(f"✗ Failed: {e}")
         
-        # Strategy 2: bf16 to CPU
-        print("\n[2/4] Trying: bf16 to CPU...")
+        # --------------------------------------------------------------
+        # Strategy 2: BF16 GPU
+        # --------------------------------------------------------------
+        print("\n[2/4] Trying: bf16 GPU...")
         try:
             model = AutoModelForCausalLM.from_pretrained(
                 config["base_model"],
                 torch_dtype=torch.bfloat16,
-                low_cpu_mem_usage=True,
-                device_map="cpu",
+                device_map="auto",
                 trust_remote_code=True,
+                low_cpu_mem_usage=True,
             )
-            print("✓ Success: bf16 to CPU")
+            print("✓ Success: bf16 GPU")
         except Exception as e:
+            errors.append(("bf16 GPU", e))
             print(f"✗ Failed: {e}")
             
-            # Strategy 3: 4-bit to CPU
-            print("\n[3/4] Trying: 4-bit to CPU...")
+            # ----------------------------------------------------------
+            # Strategy 3: BF16 CPU
+            # ----------------------------------------------------------
+            print("\n[3/4] Trying: bf16 CPU...")
             try:
                 model = AutoModelForCausalLM.from_pretrained(
                     config["base_model"],
-                    torch_dtype=torch.float16,
+                    torch_dtype=torch.bfloat16,
                     device_map="cpu",
                     trust_remote_code=True,
+                    low_cpu_mem_usage=True,
                 )
-                print("✓ Success: 4-bit to CPU")
+                print("✓ Success: bf16 CPU")
             except Exception as e:
+                errors.append(("bf16 CPU", e))
                 print(f"✗ Failed: {e}")
                 
-                # Strategy 4: 4-bit auto
-                print("\n[4/4] Trying: 4-bit auto...")
+                # ------------------------------------------------------
+                # Strategy 4: FP16 GPU
+                # ------------------------------------------------------
+                print("\n[4/4] Trying: fp16 GPU...")
                 try:
                     model = AutoModelForCausalLM.from_pretrained(
                         config["base_model"],
                         torch_dtype=torch.float16,
                         device_map="auto",
                         trust_remote_code=True,
+                        low_cpu_mem_usage=True,
                     )
-                    print("✓ Success: 4-bit auto")
+                    print("✓ Success: fp16 GPU")
                 except Exception as e:
+                    errors.append(("fp16 GPU", e))
                     print(f"✗ Failed: {e}")
-                    raise RuntimeError("All loading strategies failed!")
+                    msg = "\n".join(
+                        f"{name}: {err}" for name, err in errors
+                    )
+                    raise RuntimeError(
+                        f"All loading strategies failed:\n\n{msg}"
+                    )
     
     # Add LoRA
     lora_config = LoraConfig(