From da5eb3abed56376488be25a06ab3e0940e1de29b Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 07:45:33 -0400
Subject: [PATCH] feat: implement QLoRA with 4-bit BitsAndBytes quantization

---
 training/configs/ornith-35b-lora.yaml | 11 +----------
 training/scripts/train.py             | 19 +++++++++++++------
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml
index 4abbc07..c10fb90 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-FP8
+base_model: /data/models/Ornith-1.0-35B
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 
@@ -53,18 +53,9 @@ huggingface_hub:
 deepspeed_config:
   zero_optimization:
     stage: 3
-    offload_optimizer:
-      device: cpu
-      pin_memory: true
-    offload_param:
-      device: cpu
-      pin_memory: true
-      offload_params_device: cpu
   gradient_clipping: 1.0
   train_batch_size: auto
   train_micro_batch_size_per_gpu: auto
-  # Offload 32GB to RAM
-  zero_hierarchical_offload: true
 
 # Evaluation
 eval_strategy: steps
diff --git a/training/scripts/train.py b/training/scripts/train.py
index 0462e42..43bdc6a 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -32,17 +32,24 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model and convert to bf16 (ignore FP8 quantization)
+    # Load model with QLoRA (4-bit quantization)
     print(f"Loading model: {config['base_model']}")
+    from transformers import BitsAndBytesConfig
+    
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+    
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        dtype=torch.bfloat16,  # Convert FP8 -> bf16
-        device_map="cpu",  # Load to CPU first
+        quantization_config=quantization_config,
+        device_map="auto",  # Distribute across GPUs
         trust_remote_code=True,
     )
-    # Remove quantization config to avoid SFTTrainer validation error
-    model.config.quantization_config = None
-    print("Model loaded and converted to bf16.")
+    print("Model loaded with QLoRA (4-bit).")
     
     # Add LoRA
     lora_config = LoraConfig(