From 2fa52cd9d43695c39f24176e7beff9607caa786f Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 09:37:20 -0400
Subject: [PATCH] fix: load QLoRA model to CPU first, DeepSpeed distributes

---
 training/scripts/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 75c1cef..fb1f463 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -47,7 +47,7 @@ def train(config_path):
         config["base_model"],
         quantization_config=bnb_config,
         dtype=torch.bfloat16,
-        device_map="auto",
+        device_map="cpu",  # Load to CPU first
         trust_remote_code=True,
     )
     print("Model loaded with QLoRA (4-bit).")