diff --git a/train.py b/train.py
index c50c6d1..2117664 100644
--- a/train.py
+++ b/train.py
@@ -55,25 +55,18 @@ def train(config_path):
         print(f"✗ Failed: {e}")
         
         # --------------------------------------------------------------
-        # Strategy 2: QLoRA 4-bit with FSDP (load to CPU, FSDP shards across GPUs)
+        # Strategy 2: 4-bit model with FSDP (load to GPU, FSDP shards)
         # --------------------------------------------------------------
-        print("\n[2/4] Trying: 4-bit QLoRA (FSDP, load to CPU)...")
+        print("\n[2/4] Trying: 4-bit model with FSDP (load to GPU)...")
         try:
-            bnb_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_quant_type="nf4",
-                bnb_4bit_compute_dtype=torch.bfloat16,
-                bnb_4bit_use_double_quant=True,
-                bnb_4bit_quant_storage=torch.bfloat16,  # Enable FSDP sharding of 4-bit weights
-            )
             model = AutoModelForCausalLM.from_pretrained(
                 config["base_model"],
-                quantization_config=bnb_config,
-                device_map="cpu",  # Load to CPU, FSDP shards later
+                device_map="auto",  # Load to GPU first
+                torch_dtype=torch.float16,
                 trust_remote_code=True,
                 low_cpu_mem_usage=True,
             )
-            print("✓ Success: QLoRA 4-bit loaded to CPU (FSDP will shard across GPUs)")
+            print("✓ Success: 4-bit model loaded to GPU (FSDP will shard)")
         except Exception as e:
             errors.append(("QLoRA 4-bit FSDP", e))
             print(f"✗ Failed: {e}")
@@ -188,14 +181,14 @@ def train(config_path):
                 transformer_layer_cls={Qwen3_5MoeDecoderLayer},
             )
         
-        # Wrap model with FSDP on CPU first
-        print("Wrapping model with FSDP on CPU...")
+        # Wrap model with FSDP on GPU
+        print("Wrapping model with FSDP on GPU...")
         model = FSDP(
             model,
             auto_wrap_policy=get_auto_wrap_policy(model),
-            device_id=None,  # Keep on CPU initially
+            device_id=torch.cuda.current_device(),  # Keep on current GPU
             mixed_precision=None,
-            sync_module_states=False,  # Model is on CPU, no sync needed
+            sync_module_states=False,  # Model is already on GPU
             use_orig_params=True,
         )
         print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)")