From 815115e97abbe75396ecb831b0eba32b6052d3d5 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 21:14:31 -0400
Subject: [PATCH] feat: add BnB 4-bit quantization strategy + use bf16 model

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

diff --git a/train.py b/train.py
index a4667af..0829e88 100644
--- a/train.py
+++ b/train.py
@@ -35,44 +35,49 @@ def train(config_path):
     # Load model - try multiple strategies
     print(f"\n[INFO] Loading {config['base_model']}...")
     
-    # Strategy 1: bf16 to CPU with low memory usage
-    print("\n[1/4] Trying: bf16 to CPU...")
+    # Strategy 1: bf16 model with BnB 4-bit quantization
+    print("\n[1/4] Trying: bf16 with BnB 4-bit...")
     try:
+        from transformers import BitsAndBytesConfig
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
-            torch_dtype=torch.bfloat16,
-            low_cpu_mem_usage=True,
-            device_map="cpu",
+            quantization_config=bnb_config,
+            device_map="auto",
             trust_remote_code=True,
         )
-        print("✓ Success: bf16 to CPU")
+        print("✓ Success: bf16 with BnB 4-bit")
     except Exception as e:
         print(f"✗ Failed: {e}")
         
-        # Strategy 2: 4-bit to CPU
-        print("\n[2/4] Trying: 4-bit to CPU...")
+        # Strategy 2: bf16 to CPU
+        print("\n[2/4] Trying: bf16 to CPU...")
         try:
             model = AutoModelForCausalLM.from_pretrained(
                 config["base_model"],
-                torch_dtype=torch.float16,
+                torch_dtype=torch.bfloat16,
+                low_cpu_mem_usage=True,
                 device_map="cpu",
                 trust_remote_code=True,
             )
-            print("✓ Success: 4-bit to CPU")
+            print("✓ Success: bf16 to CPU")
         except Exception as e:
             print(f"✗ Failed: {e}")
             
-            # Strategy 3: bf16 auto
-            print("\n[3/4] Trying: bf16 auto...")
+            # Strategy 3: 4-bit to CPU
+            print("\n[3/4] Trying: 4-bit to CPU...")
             try:
                 model = AutoModelForCausalLM.from_pretrained(
                     config["base_model"],
-                    torch_dtype=torch.bfloat16,
-                    device_map="auto",
-                    low_cpu_mem_usage=True,
+                    torch_dtype=torch.float16,
+                    device_map="cpu",
                     trust_remote_code=True,
                 )
-                print("✓ Success: bf16 auto")
+                print("✓ Success: 4-bit to CPU")
             except Exception as e:
                 print(f"✗ Failed: {e}")
                 
diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml
index fedab50..718c801 100644
--- a/training/configs/ornith-35b-lora.yaml
+++ b/training/configs/ornith-35b-lora.yaml
@@ -1,7 +1,7 @@
 # LoRA Training Configuration for Llama-2-7b
 # 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: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer