From 04a95f362f9af07339b7d037c31bb1c8f23b828a Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 21:09:54 -0400
Subject: [PATCH] fix: load bf16 model directly with BnB 4-bit (single load)

---
 train.py                              | 23 +++--------------------
 training/configs/ornith-35b-lora.yaml |  2 +-
 2 files changed, 4 insertions(+), 21 deletions(-)

diff --git a/train.py b/train.py
index a547fb8..b261970 100644
--- a/train.py
+++ b/train.py
@@ -33,19 +33,8 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load bf16 to CPU first
-    print(f"\n[INFO] Loading {config['base_model']} bf16 to CPU...")
-    model = AutoModelForCausalLM.from_pretrained(
-        config["base_model"],
-        device_map="cpu",
-        torch_dtype=torch.bfloat16,
-        trust_remote_code=True,
-        low_cpu_mem_usage=True,
-    )
-    print("✓ Model loaded to CPU (~70GB bf16)")
-    
-    # Quantize with BnB on CPU
-    print("   Quantizing with BnB 4-bit on CPU...")
+    # Load bf16 to CPU with BnB 4-bit quantization
+    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,
@@ -53,12 +42,6 @@ def train(config_path):
         bnb_4bit_compute_dtype=torch.float16,
     )
     
-    # Reload with quantization
-    del model
-    import gc
-    gc.collect()
-    torch.cuda.empty_cache()
-    
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
         quantization_config=bnb_config,
@@ -67,7 +50,7 @@ def train(config_path):
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
-    print("✓ Model quantized to 4-bit on CPU")
+    print("✓ Model loaded to CPU with BnB 4-bit (~17.5GB)")
     
     # Move to GPU
     print("   Moving to GPU 0...")
diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml
index 32a05b4..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-bnb-4bit
+base_model: /data/models/Ornith-1.0-35B
 model_type: Qwen3_5MoeForCausalLM
 tokenizer_type: AutoTokenizer