[PPDiffusers] update ldm train code (#3772)

* update ldm train

* add rank==0

* add rank==0

* update trainer

* update

* update

* update

* update

* update trainer code

* update generate pipelines

* update support bf16

Author: JunnYu

paddlenlp/trainer/trainer.py

@@ -177,7 +177,7 @@ class Trainer:
     def __init__(
         self,
         model: Union[PretrainedModel, nn.Layer] = None,
-        criterion: Union[nn.Layer] = None,
+        criterion: nn.Layer = None,
         args: TrainingArguments = None,
         data_collator: Optional[DataCollator] = None,
         train_dataset: Optional[Dataset] = None,
@@ -241,6 +241,7 @@ def __init__(
         self.state = TrainerState()
         self.control = TrainerControl()
         self._signature_columns = None
+        self.optimizer_grouped_parameters = None
 
         if (self.sharding is not None) and (self.optimizer is not None
                                             or self.lr_scheduler is not None):
@@ -710,9 +711,11 @@ def train(
 
                     self.control = self.callback_handler.on_step_end(
                         args, self.state, self.control)
-
-                    self._maybe_log_save_evaluate(tr_loss, model, epoch,
-                                                  ignore_keys_for_eval)
+                    self._maybe_log_save_evaluate(tr_loss,
+                                                  model,
+                                                  epoch,
+                                                  ignore_keys_for_eval,
+                                                  inputs=inputs)
                 else:
                     self.control = self.callback_handler.on_substep_end(
                         args, self.state, self.control)
@@ -730,8 +733,11 @@ def train(
 
             self.control = self.callback_handler.on_epoch_end(
                 args, self.state, self.control)
-            self._maybe_log_save_evaluate(tr_loss, model, epoch,
-                                          ignore_keys_for_eval)
+            self._maybe_log_save_evaluate(tr_loss,
+                                          model,
+                                          epoch,
+                                          ignore_keys_for_eval,
+                                          inputs=inputs)
 
             if self.control.should_training_stop:
                 break
@@ -805,7 +811,7 @@ def _set_state_dict_in_model(self, state_dict):
         self.model.set_state_dict(state_dict)
 
     def _maybe_log_save_evaluate(self, tr_loss, model, epoch,
-                                 ignore_keys_for_eval):
+                                 ignore_keys_for_eval, **kwargs):
         if self.control.should_log:
 
             logs: Dict[str, float] = {}
@@ -836,7 +842,7 @@ def _maybe_log_save_evaluate(self, tr_loss, model, epoch,
             self._globalstep_last_logged = self.state.global_step
             self._globalstep_last_start_time = time.time()
 
-            self.log(logs)
+            self.log(logs, **kwargs)
 
         metrics = None
         if self.control.should_evaluate:
@@ -1024,11 +1030,16 @@ def create_optimizer(self, lr_scheduler=None):
         Trainer's init through `optimizers`, or subclass and override this method in a subclass.
         """
         if self.optimizer is None:
-            decay_parameters = [
-                p.name for n, p in self.model.named_parameters()
-                if not any(nd in n for nd in ["bias", "norm"])
-            ]
-            apply_decay_param_fun = lambda x: x in decay_parameters
+            if self.optimizer_grouped_parameters is not None:
+                params = self.optimizer_grouped_parameters
+                apply_decay_param_fun = None
+            else:
+                params = self.model.parameters()
+                decay_parameters = [
+                    p.name for n, p in self.model.named_parameters()
+                    if not any(nd in n for nd in ["bias", "norm"])
+                ]
+                apply_decay_param_fun = lambda x: x in decay_parameters
 
             optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
                 self.args)
@@ -1038,22 +1049,24 @@ def create_optimizer(self, lr_scheduler=None):
                 self.optimizer = DygraphShardingOptimizer(
                     hcg=fleet.get_hybrid_communicate_group(),
                     user_defined_strategy=None,
-                    params=self.model.parameters(),
+                    params=params,
                     inner_optimizer_class=optimizer_cls,
                     learning_rate=self.lr_scheduler
                     if lr_scheduler is None else lr_scheduler,
                     apply_decay_param_fun=apply_decay_param_fun,
                     weight_decay=self.args.weight_decay,
-                    grad_clip=nn.ClipGradByGlobalNorm(self.args.max_grad_norm),
+                    grad_clip=nn.ClipGradByGlobalNorm(self.args.max_grad_norm)
+                    if self.args.max_grad_norm > 0 else None,
                     **optimizer_kwargs)
             else:
                 self.optimizer = optimizer_cls(
                     learning_rate=self.lr_scheduler
                     if lr_scheduler is None else lr_scheduler,
                     apply_decay_param_fun=apply_decay_param_fun,
-                    parameters=self.model.parameters(),
+                    parameters=params,
                     weight_decay=self.args.weight_decay,
-                    grad_clip=nn.ClipGradByGlobalNorm(self.args.max_grad_norm),
+                    grad_clip=nn.ClipGradByGlobalNorm(self.args.max_grad_norm)
+                    if self.args.max_grad_norm > 0 else None,
                     **optimizer_kwargs)
 
         return self.optimizer
@@ -1429,6 +1442,10 @@ def _save_checkpoint(self, model, metrics=None):
         if self.args.should_save:
             self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)
 
+    def set_optimizer_grouped_parameters(self,
+                                         optimizer_grouped_parameters=None):
+        self.optimizer_grouped_parameters = optimizer_grouped_parameters
+
     def _sorted_checkpoints(self,
                             output_dir=None,
                             checkpoint_prefix=PREFIX_CHECKPOINT_DIR,
@@ -1553,7 +1570,7 @@ def _load_optimizer_and_scheduler(self, checkpoint):
                     paddle.load(os.path.join(checkpoint, SCALER_NAME),
                                 return_numpy=True))
 
-    def log(self, logs: Dict[str, float]) -> None:
+    def log(self, logs: Dict[str, float], **kwargs) -> None:
         """
         Log `logs` on the various objects watching training.
 
@@ -1569,7 +1586,8 @@ def log(self, logs: Dict[str, float]) -> None:
         output = {**logs, **{"step": self.state.global_step}}
         self.state.log_history.append(output)
         self.control = self.callback_handler.on_log(self.args, self.state,
-                                                    self.control, logs)
+                                                    self.control, logs,
+                                                    **kwargs)
 
     def evaluate(
         self,

paddlenlp/trainer/trainer_callback.py

@@ -410,9 +410,14 @@ def on_save(self, args: TrainingArguments, state: TrainerState,
         return self.call_event("on_save", args, state, control)
 
     def on_log(self, args: TrainingArguments, state: TrainerState,
-               control: TrainerControl, logs):
+               control: TrainerControl, logs, **kwargs):
         control.should_log = False
-        return self.call_event("on_log", args, state, control, logs=logs)
+        return self.call_event("on_log",
+                               args,
+                               state,
+                               control,
+                               logs=logs,
+                               **kwargs)
 
     def on_prediction_step(self, args: TrainingArguments, state: TrainerState,
                            control: TrainerControl):