From 2fa5173c561d796d46b5acf88badb8191b46a8af Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 07:48:44 -0400
Subject: [PATCH] feat: add prepare_model_for_kbit_training and update LoRA
 r=64 alpha=128

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

diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml
index c10fb90..73fd9fb 100644
--- a/training/configs/ornith-35b-lora.yaml
+++ b/training/configs/ornith-35b-lora.yaml
@@ -9,8 +9,8 @@ tokenizer_type: LlamaTokenizer
 # No need for BitsAndBytes configuration
 
 # LoRA Configuration
-lora_r: 16
-lora_alpha: 32
+lora_r: 64
+lora_alpha: 128
 lora_dropout: 0.05
 target_modules:
   - q_proj
diff --git a/training/scripts/train.py b/training/scripts/train.py
index 0bcbe3b..75c1cef 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -52,6 +52,11 @@ def train(config_path):
     )
     print("Model loaded with QLoRA (4-bit).")
     
+    # Prepare model for k-bit training
+    from peft import prepare_model_for_kbit_training
+    model = prepare_model_for_kbit_training(model)
+    print("Model prepared for k-bit training.")
+    
     # Add LoRA
     lora_config = LoraConfig(
         r=config["lora_r"],