From af09d47c6f8c1dcbe29f3f03cfea81cb899049e0 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 00:03:24 -0400
Subject: [PATCH] fix: load model to CPU first, let DeepSpeed distribute

---
 training/scripts/train.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 4e92eae..029a2a7 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -32,13 +32,14 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model - preserve NVFP4 quantization
-    # Use device_map="auto" to distribute across GPUs from the start
+    # Load model to CPU first, DeepSpeed will distribute
+    print("Loading model to CPU...")
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        device_map="auto",  # Distribute across GPUs
+        device_map="cpu",  # Load to CPU first
         trust_remote_code=True,
     )
+    print(f"Model loaded to CPU. Moving to GPU with DeepSpeed...")
     
     # Add LoRA
     lora_config = LoraConfig(