From 1c782b58139272301ea288c4d8d1ae1f96d23a16 Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Mon, 4 Dec 2023 12:40:43 -0500 Subject: [PATCH] in peft finetune, only the trainable parameters need to be saved to reduce the storage size and also save the time of checkpoint saving while using deepspeed for training Signed-off-by: Wang, Yi --- src/transformers/trainer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 3c9e4420124012..aaec8d23db8b71 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2419,7 +2419,13 @@ def _save_optimizer_and_scheduler(self, output_dir): elif self.is_deepspeed_enabled: # under zero3 model file itself doesn't get saved since it's bogus! Unless deepspeed # config `stage3_gather_16bit_weights_on_model_save` is True - self.model_wrapped.save_checkpoint(output_dir) + accept_exclude_frozen_parameters = "exclude_frozen_parameters" in set( + inspect.signature(self.model_wrapped.save_checkpoint).parameters.keys() + ) + if accept_exclude_frozen_parameters and is_peft_available() and isinstance(self.model, PeftModel): + self.model_wrapped.save_checkpoint(output_dir, exclude_frozen_parameters=True) + else: + self.model_wrapped.save_checkpoint(output_dir) elif self.is_fsdp_enabled: # save fsdp specific ckpt for resuming from ckpt save_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, self.model, output_dir)