From 03a0f236a483d2b05094cc6e5546cfa28bd1668b Mon Sep 17 00:00:00 2001 From: FP HAM Date: Wed, 8 Jan 2025 14:54:09 -0500 Subject: [PATCH] Training_PRO fix: add `if 'quantization_config' in shared.model.config.to_dict()` --- extensions/Training_PRO/script.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/extensions/Training_PRO/script.py b/extensions/Training_PRO/script.py index 5365154c..01bcf67d 100644 --- a/extensions/Training_PRO/script.py +++ b/extensions/Training_PRO/script.py @@ -789,7 +789,11 @@ def do_train(lora_name: str, always_override: bool, save_steps: int, micro_batch if not hasattr(shared.model, 'lm_head') or hasattr(shared.model.lm_head, 'weight'): logger.info("Getting model ready...") # here we can disable gradient checkpoint, by default = true, use_gradient_checkpointing=True - prepare_model_for_kbit_training(shared.model) + if 'quantization_config' in shared.model.config.to_dict(): + print(f"Method: {RED}QLORA{RESET}") + prepare_model_for_kbit_training(shared.model) + else: + print(f"Method: {RED}LoRA{RESET}") # base model is now frozen and should not be reused for any other LoRA training than this one shared.model_dirty_from_training = True