feat: add prepare_model_for_kbit_training and update LoRA r=64 alpha=128

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

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