diff --git a/Context.py b/Context.py index e226590..cd9e131 100644 --- a/Context.py +++ b/Context.py @@ -1,9 +1,12 @@ import configparser import torch +from peft import PeftModel from sympy.printing import str class Context: + model: PeftModel + accel: bool device: str purgeTargetDirectories: bool diff --git a/Merger.py b/Merger.py index cd2d2c5..b6933a6 100644 --- a/Merger.py +++ b/Merger.py @@ -7,15 +7,17 @@ class Merger: def mergeAndStore(self, context: Context): if not context.mergeFull: return - if not context.storeAdapter: - raise Exception("Configuration problem: mergeFull requires storeAdapter") - print("Reloading model") - baseModel = ModelLoader().load(context.locBaseModel, False, True, context) - peftModel = PeftModel.from_pretrained(baseModel, context.locAdapter) + if not context.model: + if not context.storeAdapter: + raise Exception("Configuration problem: mergeFull requires storeAdapter") + + print("Loading model") + baseModel = ModelLoader().load(context.locBaseModel, False, True, context) + context.model = PeftModel.from_pretrained(baseModel, context.locAdapter) print("Merging") - mergedModel = peftModel.merge_and_unload() + mergedModel = context.model.merge_and_unload() print("Dequantize") try: diff --git a/Trainer.py b/Trainer.py index 8b085b4..6a8363e 100644 --- a/Trainer.py +++ b/Trainer.py @@ -16,42 +16,36 @@ class Trainer: baseModel = ModelLoader().load(context.locBaseModel, False, True, context) sftTrainer = self.__createTrainer(baseModel, context) - try: - vInplace = context.validate and context.vInplace - validator = None - if vInplace: - validator = Validator() - - try: - cnt = 0 - continueTraining = True - while continueTraining: - cnt += 1 - print(f"Training run {cnt}") - - self.__train(sftTrainer) - - if vInplace: - continueTraining = not validator.validateInPlace(sftTrainer.tokenizer, sftTrainer.model, context) - else: - continueTraining = False - - print(f"Total training runs: {cnt} epochs {cnt * (context.trEpochs or 1)}") - - finally: - if vInplace: - print(f"Validation statistics:\n{validator.statistics}") - validator.unload(context) - - if context.storeAdapter: - self.__storeAdapter(sftTrainer, context) + context.model = sftTrainer.model + vInplace = context.validate and context.vInplace + validator = None + if vInplace: + validator = Validator() + + try: + cnt = 0 + continueTraining = True + while continueTraining: + cnt += 1 + print(f"Training run {cnt}") + + self.__train(sftTrainer) + + if vInplace: + continueTraining = not validator.validateInPlace(sftTrainer.tokenizer, sftTrainer.model, context) + else: + continueTraining = False + + print(f"Total training runs: {cnt} epochs {cnt * (context.trEpochs or 1)}") + finally: - del sftTrainer - del baseModel - gc.collect() - if context.accel: - torch.cuda.empty_cache() + if vInplace: + print(f"Validation statistics:\n{validator.statistics}") + validator.unload(context) + + if context.storeAdapter: + self.__storeAdapter(sftTrainer, context) def __createTrainer(self, baseModel, context: Context): tokenizer = AutoTokenizer.from_pretrained(