fix the bug

Author: zty-king

paddlenlp/trainer/trainer.py

@@ -2855,7 +2855,7 @@ def _save_checkpoint(self, model, metrics=None):
             if (
                 self.args.should_save
                 or self.args.use_expert_parallel
-                or (self.args.data_parallel_degree > 1 and not self.args.use_hybrid_parallel)
+                or (self.args.data_parallel_degree > 1 and self.args.save_flex_checkpoint)
             ):
                 if not self.args.use_hybrid_parallel:
                     logger.info("Saving optimizer files.")