diff --git a/train.py b/train.py
index 6991b6c..c88a33c 100644
--- a/train.py
+++ b/train.py
@@ -32,63 +32,22 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model - try multiple strategies
-    print(f"\n[INFO] Loading {config['base_model']}...")
+    # Load model with memory optimization
+    print(f"\n[INFO] Loading {config['base_model']} with memory optimization...")
     
-    # Strategy 1: 4-bit AS-IS (already quantized)
-    print("\n[1/4] Trying: 4-bit AS-IS...")
+    # Direct low-memory load for large model + DeepSpeed
     try:
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
-            torch_dtype=torch.float16,
-            # device_map="auto",
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            device_map="cpu",           # Force CPU first
             trust_remote_code=True,
         )
-        print("✓ Success: 4-bit AS-IS")
+        print("✓ Success: Loaded on CPU with low memory usage")
     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!")
+        print(f"✗ Failed initial load: {e}")
+        raise
     
     # Add LoRA
     lora_config = LoraConfig(