diff --git a/training/scripts/train.py b/training/scripts/train.py
index a432f72..95fef47 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -59,70 +59,42 @@ def train(config_path):
     except Exception as e:
         print(f"✗ Failed: {e}")
         
-        # Strategy 2: Load bf16 to CPU with low_cpu_mem_usage
-        print("\n[2/4] Trying: bf16 model to CPU...")
+        # Strategy 2: DeepSpeed ZeRO-3 (distribute across GPUs)
+        print("\n[2/4] Trying: DeepSpeed ZeRO-3...")
         try:
+            import deepspeed
             model = AutoModelForCausalLM.from_pretrained(
                 config["base_model"],
                 torch_dtype=torch.bfloat16,
-                device_map="cpu",
+                device_map=None,
                 low_cpu_mem_usage=True,
                 trust_remote_code=True,
             )
-            print("✓ Success: bf16 model loaded to CPU")
+            
+            ds_config = {
+                "train_micro_batch_size_per_gpu": 1,
+                "gradient_accumulation_steps": 1,
+                "zero_optimization": {
+                    "stage": 3,
+                    "contiguous_gradients": True,
+                    "overlap_comm": True,
+                    "offload_optimizer": {"device": "cpu", "pin_memory": True},
+                    "offload_param": {"device": "cpu", "pin_memory": True},
+                },
+                "bf16": {"enabled": True},
+            }
+            
+            optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train_params"]["learning_rate"]))
+            model, optimizer, _, _ = deepspeed.initialize(
+                model=model,
+                model_parameters=model.parameters(),
+                optimizer=optimizer,
+                config=ds_config,
+            )
+            print("✓ Success: DeepSpeed ZeRO-3 loaded")
         except Exception as e:
             print(f"✗ Failed: {e}")
-            
-            # Strategy 3: Load fp16 with auto placement
-            print("\n[3/4] Trying: fp16 model with auto placement...")
-            try:
-                model = AutoModelForCausalLM.from_pretrained(
-                    config["base_model"],
-                    torch_dtype=torch.float16,
-                    device_map="auto",
-                    low_cpu_mem_usage=True,
-                    trust_remote_code=True,
-                )
-                print("✓ Success: fp16 model loaded")
-            except Exception as e:
-                print(f"✗ Failed: {e}")
-                
-                # Strategy 4: DeepSpeed ZeRO-3 with CPU offload
-                print("\n[4/4] Trying: DeepSpeed ZeRO-3 CPU offload...")
-                try:
-                    import deepspeed
-                    model = AutoModelForCausalLM.from_pretrained(
-                        config["base_model"],
-                        torch_dtype=torch.bfloat16,
-                        device_map=None,
-                        low_cpu_mem_usage=True,
-                        trust_remote_code=True,
-                    )
-                    
-                    ds_config = {
-                        "train_micro_batch_size_per_gpu": 1,
-                        "gradient_accumulation_steps": 1,
-                        "zero_optimization": {
-                            "stage": 3,
-                            "contiguous_gradients": True,
-                            "overlap_comm": True,
-                            "offload_optimizer": {"device": "cpu", "pin_memory": True},
-                            "offload_param": {"device": "cpu", "pin_memory": True},
-                        },
-                        "bf16": {"enabled": True},
-                    }
-                    
-                    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
-                    model, optimizer, _, _ = deepspeed.initialize(
-                        model=model,
-                        model_parameters=model.parameters(),
-                        optimizer=optimizer,
-                        config=ds_config,
-                    )
-                    print("✓ Success: DeepSpeed ZeRO-3 loaded")
-                except Exception as e:
-                    print(f"✗ Failed: {e}")
-                    raise RuntimeError("All loading strategies failed!")
+            raise RuntimeError("All loading strategies failed!")
     
     # Prepare model for k-bit training
     from peft import prepare_model_for_kbit_training