From a0f4f644b018d2034c104fcf9141097dee56873c Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 19:33:37 -0400
Subject: [PATCH] fix: load 4-bit model with device_map=auto (let transformers
 distribute)

---
 train.py | 53 ++++++++++-------------------------------------------
 1 file changed, 10 insertions(+), 43 deletions(-)

diff --git a/train.py b/train.py
index 4ee0b92..024ddf6 100644
--- a/train.py
+++ b/train.py
@@ -33,55 +33,22 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load bf16 model to CPU
-    print(f"\n[INFO] Loading {config['base_model']} bf16 to CPU...")
+    # Load 4-bit model directly (already quantized on disk)
+    print(f"\n[INFO] Loading {config['base_model']} (4-bit)...")
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        device_map="cpu",
-        torch_dtype=torch.bfloat16,
+        device_map="auto",  # Let transformers distribute
+        torch_dtype=torch.float16,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
-    print("✓ Model loaded to CPU (~70GB bf16)")
+    print("✓ Success: Model loaded")
     
-    # Apply PEFT k-bit training preparation
-    print("   Applying PEFT k-bit preparation...")
-    from peft import prepare_model_for_kbit_training
-    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
-    print("   ✓ Model prepared for k-bit training")
-    
-    # Manually quantize linear layers
-    print("   Quantizing linear layers to 4-bit...")
-    from bitsandbytes.nn import Linear4bit
-    from torch import nn
-    
-    quantized_count = 0
-    for name, module in model.named_modules():
-        if isinstance(module, nn.Linear) and 'lm_head' not in name:
-            new_module = Linear4bit(
-                module.in_features,
-                module.out_features,
-                bias=module.bias is not None,
-            )
-            new_module.weight = nn.Parameter(module.weight.data.clone())
-            if module.bias is not None:
-                new_module.bias = nn.Parameter(module.bias.data.clone())
-            
-            layers = name.split('.')
-            parent = model
-            for layer in layers[:-1]:
-                parent = getattr(parent, layer)
-            setattr(parent, layers[-1], new_module)
-            quantized_count += 1
-    
-    print(f"   ✓ Quantized {quantized_count} linear layers to 4-bit")
-    
-    # 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")
+    # Check memory usage
+    for i in range(torch.cuda.device_count()):
+        mem = torch.cuda.memory_allocated(i) / 1e9
+        total = torch.cuda.get_device_properties(i).total_memory / 1e9
+        print(f"   GPU {i}: {mem:.2f} GB / {total:.2f} GB")
     
     # Add LoRA
     lora_config = LoraConfig(