[unified checkpoint] Fix last checkpoint save

paddlenlp/trainer/trainer.py

@@ -2027,8 +2027,17 @@ def save_model(self, output_dir: Optional[str] = None, merge_tensor_parallel: Op
             self.model_wrapped.get_all_parameters(convert2cpu=True)
 
         if self.args.should_save_model_state:
+            unified_checkpoint_config_backup = self.args.unified_checkpoint_config
+            # backup and remove unified_checkpoint_config for not trine stage
+            if not self.is_in_train:
+                self.args.unified_checkpoint_config = []
+
             self._save(output_dir=output_dir, merge_tensor_parallel=merge_tensor_parallel)
 
+            # recover unified_checkpoint_config for not trine stage
+            if not self.is_in_train:
+                self.args.unified_checkpoint_config = unified_checkpoint_config_backup
+
     def _save_checkpoint(self, model, metrics=None):
         # assert unwrap_model(model) is self.model, "internal model should be a reference to self.model"
 

paddlenlp/trainer/training_args.py

@@ -1310,6 +1310,8 @@ def is_segment_parallel_supported():
                     "master_weight_compatible",
                     "async_save",
                 ]
+            else:
+                self.unified_checkpoint_config = self.unified_checkpoint_config.split(" ")
 
         if self.report_to is None:
             logger.info(