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"],