diff --git a/train.py b/train.py
index 54a57fa..c7175a4 100644
--- a/train.py
+++ b/train.py
@@ -33,48 +33,18 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model with accelerate device_map (Test 7 method - DISTRIBUTED)
-    print(f"\n[INFO] Loading {config['base_model']}...")
-    
-    # Detect layer names dynamically
-    from transformers import AutoConfig
-    print("   Detecting layer names...")
-    config_model = AutoConfig.from_pretrained(config["base_model"], trust_remote_code=True)
-    layer_names = [f"{config_model.model_type.title().replace('_', '')}DecoderLayer"]
-    if 'moe' in config_model.model_type.lower():
-        layer_names.append(f"{config_model.model_type.title().replace('_', '')}SparseMoeBlock")
-    print(f"   Detected layers: {layer_names}")
-    
-    # Use accelerate to create device_map
-    from accelerate import infer_auto_device_map
-    print("   Creating device_map with accelerate...")
-    
-    # First load to CPU to get the model structure
-    model_temp = AutoModelForCausalLM.from_pretrained(
-        config["base_model"],
-        device_map="cpu",
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-        low_cpu_mem_usage=True,
-    )
-    
-    # Create device_map
-    device_map = infer_auto_device_map(
-        model_temp,
-        max_memory={i: "15GB" for i in range(torch.cuda.device_count())},
-        no_split_module_classes=layer_names,
-    )
-    print(f"   Created device_map with {len(device_map)} entries")
-    
-    # Reload with device_map
+    # Load model to single GPU (MoE needs all layers on same device)
+    print(f"\n[INFO] Loading {config['base_model']} to single GPU...")
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        device_map=device_map,
+        device_map="cuda:0",
         torch_dtype=torch.float16,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
-    print("✓ Success: Model distributed across GPUs via accelerate device_map")
+    print("✓ Success: Model loaded to GPU 0")
+    print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    print(f"   Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
     
     # Add LoRA
     lora_config = LoraConfig(
@@ -107,12 +77,10 @@ def train(config_path):
         },
     )
     
-    # Model is already distributed across GPUs via device_map
-    # No FSDP needed - device_map handles distribution
-    print("✓ Model already distributed across GPUs (device_map)")
+    # Model is on single GPU
+    print("✓ Model loaded to single GPU")
     print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
-    if torch.cuda.device_count() > 1:
-        print(f"   GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
+    print(f"   Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
     
     # Training arguments
     training_args = TrainingArguments(