From c607888d8fbc2fd8abb86d4c6ae869761c3e11fc Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 08:26:23 -0400
Subject: [PATCH] fix: load QLoRA model to CPU first, FSDP shards later

---
 train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index c7d3175..a6ef3ce 100644
--- a/train.py
+++ b/train.py
@@ -50,11 +50,11 @@ def train(config_path):
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
             quantization_config=bnb_config,
-            device_map="auto",
+            device_map="cpu",  # Load to CPU, FSDP shards later
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
-        print("✓ Success: QLoRA 4-bit")
+        print("✓ Success: QLoRA 4-bit loaded to CPU")
     except Exception as e:
         errors.append(("QLoRA 4-bit", e))
         print(f"✗ Failed: {e}")