From 2d13e63c256db94e031b80be12860628296ea237 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:55:23 -0400
Subject: [PATCH] feat: add 4 loading strategies with automatic fallback

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

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 1c83c52..016717e 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -36,15 +36,68 @@ def train(config_path):
     print(f"Loading model: {config['base_model']}")
     from transformers import BitsAndBytesConfig, AutoConfig
     
-    # Load 4-bit model AS-IS (don't apply additional quantization)
+    # Try multiple loading strategies in order
     print(f"Loading model: {config['base_model']}")
-    model = AutoModelForCausalLM.from_pretrained(
-        config["base_model"],
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-    print("Model loaded (4-bit). DeepSpeed will distribute.")
+    
+    # Strategy 1: Load 4-bit AS-IS
+    print("\n[1/4] Trying: 4-bit model 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 model loaded")
+    except Exception as e:
+        print(f"✗ Failed: {e}")
+        
+        # Strategy 2: Load bf16 to CPU
+        print("\n[2/4] Trying: bf16 model to CPU...")
+        try:
+            model = AutoModelForCausalLM.from_pretrained(
+                config["base_model"],
+                torch_dtype=torch.bfloat16,
+                device_map="cpu",
+                trust_remote_code=True,
+            )
+            print("✓ Success: bf16 model loaded to CPU")
+        except Exception as e:
+            print(f"✗ Failed: {e}")
+            
+            # Strategy 3: Load with accelerate CPU offload
+            print("\n[3/4] Trying: bf16 with accelerate CPU offload...")
+            try:
+                from accelerate import load_checkpoint_and_dispatch
+                base_model = AutoModelForCausalLM.from_pretrained(
+                    config["base_model"],
+                    torch_dtype=torch.bfloat16,
+                    trust_remote_code=True,
+                )
+                model = load_checkpoint_and_dispatch(
+                    base_model,
+                    checkpoint=config["base_model"],
+                    device_map="auto",
+                    dtype=torch.bfloat16,
+                    offload_folder="/tmp/model_offload",
+                )
+                print("✓ Success: bf16 with accelerate offload")
+            except Exception as e:
+                print(f"✗ Failed: {e}")
+                
+                # Strategy 4: Use bf16 with DeepSpeed ZeRO-3
+                print("\n[4/4] Trying: bf16 with DeepSpeed ZeRO-3 CPU offload...")
+                try:
+                    model = AutoModelForCausalLM.from_pretrained(
+                        config["base_model"],
+                        torch_dtype=torch.bfloat16,
+                        device_map="cpu",
+                        trust_remote_code=True,
+                    )
+                    print("✓ Success: bf16 model (DeepSpeed will handle offload)")
+                except Exception as e:
+                    print(f"✗ Failed: {e}")
+                    raise RuntimeError("All loading strategies failed!")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training