From 42c25962b2cb60f769a3f3b0acf7dd918a3d26e0 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 18:23:39 -0400
Subject: [PATCH] feat: load model to CPU with BnB 4-bit, then move to GPU

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

diff --git a/train.py b/train.py
index 1ed40c1..9b8a033 100644
--- a/train.py
+++ b/train.py
@@ -33,16 +33,28 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model (already 4-bit on disk, just move to GPU)
-    print(f"\n[INFO] Loading {config['base_model']} to GPU 0...")
+    # Load model with BnB 4-bit on CPU, then move to GPU
+    print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit to CPU...")
+    from transformers import BitsAndBytesConfig
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+    )
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        device_map="cuda:0",
+        device_map="cpu",
+        quantization_config=bnb_config,
         torch_dtype=torch.float16,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
-    print("✓ Success: Model loaded to GPU 0")
+    print(f"   CPU RAM: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB (model on CPU)")
+    
+    # Move to GPU
+    print("   Moving to GPU 0...")
+    model = model.to("cuda:0")
+    print("✓ Success: Model loaded to GPU 0 (4-bit)")
     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")