bug fix

PaddlePaddle · Oct 17, 2024 · 5b20bd3 · 5b20bd3
1 parent 5244e61
commit 5b20bd3
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/paddlenlp/trainer/trainer.py b/paddlenlp/trainer/trainer.py
@@ -2369,12 +2369,21 @@ def _save_checkpoint(self, model, metrics=None):
         output_dir = os.path.join(run_dir, checkpoint_folder)
         signal_dir = os.path.join(run_signal_dir, checkpoint_folder)
 
-        if isinstance(self.model, LoRAModel) and (self.model.quantized or self.args.pipeline_parallel_degree > 1):
-            self.save_model(output_dir, False, signal_dir)
-        elif isinstance(self.model, LoRAModel) or isinstance(self.model, PrefixModelForCausalLM):
-            self.save_model(output_dir, True, signal_dir)
+        signature = inspect.signature(self.save_model)
+        if "signal_dir" not in signature.parameters:
+            if isinstance(self.model, LoRAModel) and (self.model.quantized or self.args.pipeline_parallel_degree > 1):
+                self.save_model(output_dir)
+            elif isinstance(self.model, LoRAModel) or isinstance(self.model, PrefixModelForCausalLM):
+                self.save_model(output_dir, True)
+            else:
+                self.save_model(output_dir)
         else:
-            self.save_model(output_dir, False, signal_dir)
+            if isinstance(self.model, LoRAModel) and (self.model.quantized or self.args.pipeline_parallel_degree > 1):
+                self.save_model(output_dir, False, signal_dir)
+            elif isinstance(self.model, LoRAModel) or isinstance(self.model, PrefixModelForCausalLM):
+                self.save_model(output_dir, True, signal_dir)
+            else:
+                self.save_model(output_dir, False, signal_dir)
 
         # only save model state dict, ignore optimizer and scheduler
         if not self.args.ignore_save_lr_and_optim: