fix: restore all strategies with CPU-first as primary

train.py

@@ -32,22 +32,63 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model with memory optimization
-    print(f"\n[INFO] Loading {config['base_model']} with memory optimization...")
+    # Load model - try multiple strategies
+    print(f"\n[INFO] Loading {config['base_model']}...")
     
-    # Direct low-memory load for large model + DeepSpeed
+    # Strategy 1: bf16 to CPU with low memory usage
+    print("\n[1/4] Trying: bf16 to CPU...")
     try:
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
             torch_dtype=torch.bfloat16,
             low_cpu_mem_usage=True,
-            device_map="cpu",           # Force CPU first
+            device_map="cpu",
             trust_remote_code=True,
         )
-        print("✓ Success: Loaded on CPU with low memory usage")
+        print("✓ Success: bf16 to CPU")
     except Exception as e:
-        print(f"✗ Failed initial load: {e}")
-        raise
+        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 auto
+            print("\n[3/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}")
+                
+                # Strategy 4: 4-bit auto
+                print("\n[4/4] Trying: 4-bit auto...")
+                try:
+                    model = AutoModelForCausalLM.from_pretrained(
+                        config["base_model"],
+                        torch_dtype=torch.float16,
+                        device_map="auto",
+                        trust_remote_code=True,
+                    )
+                    print("✓ Success: 4-bit auto")
+                except Exception as e:
+                    print(f"✗ Failed: {e}")
+                    raise RuntimeError("All loading strategies failed!")
     
     # Add LoRA
     lora_config = LoraConfig(