From 3c31e4dfb0d264b7b40a686447fc9ce4ac2ee4ec Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 19:55:01 -0400
Subject: [PATCH] feat: load bf16 model, quantize with BnB on CPU, move to GPU

---
 train.py                              | 45 +++++++++++++++++++++------
 training/configs/ornith-35b-lora.yaml |  2 +-
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/train.py b/train.py
index 024ddf6..1660f36 100644
--- a/train.py
+++ b/train.py
@@ -33,22 +33,49 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load 4-bit model directly (already quantized on disk)
-    print(f"\n[INFO] Loading {config['base_model']} (4-bit)...")
+    # Load bf16 model to CPU, then quantize with BnB
+    print(f"\n[INFO] Loading {config['base_model']} bf16 to CPU...")
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        device_map="auto",  # Let transformers distribute
+        device_map="cpu",
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+    )
+    print("✓ Model loaded to CPU (~70GB bf16)")
+    
+    # Apply BnB 4-bit quantization on CPU
+    print("   Applying BnB 4-bit quantization on CPU...")
+    from transformers import BitsAndBytesConfig
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_use_double_quant=True,
+    )
+    
+    # Reload with quantization config
+    del model
+    import gc
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        config["base_model"],
+        device_map="cpu",
+        quantization_config=bnb_config,
         torch_dtype=torch.float16,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
-    print("✓ Success: Model loaded")
+    print("✓ Model quantized to 4-bit on CPU (~17.5GB)")
     
-    # 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")
+    # 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")
     
     # Add LoRA
     lora_config = LoraConfig(
diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml
index 5088306..2cf394e 100644
--- a/training/configs/ornith-35b-lora.yaml
+++ b/training/configs/ornith-35b-lora.yaml
@@ -1,7 +1,7 @@
 # LoRA Training Configuration for Ornith-1.0-35B
 # Dataset: cyron_summary_lora_dataset (20k examples)
 
-base_model: /data/models/Ornith-1.0-35B-4bit
+base_model: /data/models/Ornith-1.0-35B
 model_type: Qwen3_5MoeForCausalLM
 tokenizer_type: AutoTokenizer