[Cherry-pick]Support to load sharded EMA checkpoint (#11075)

* support load sharded EMA checkpoints

* support_ema_loading_no_pdopt

* polish code

Author: sneaxiy

paddlenlp/trainer/trainer.py

@@ -378,6 +378,7 @@ def __init__(
                 self.model,
                 self.optimizer,
                 remap_parameter_name=self.args.load_sharded_model_remap_parameter_name,
+                is_ema=self.args.sharded_model_from_ema,
             )
 
         if self.args.unified_checkpoint:

paddlenlp/trainer/training_args.py

@@ -638,6 +638,11 @@ class TrainingArguments:
         metadata={"help": "Whether to remap parameter name when load_sharded_model = true."},
     )
 
+    sharded_model_from_ema: bool = field(
+        default=False,
+        metadata={"help": "Whether to load sharded model from EMA."},
+    )
+
     tensor_parallel_degree: int = field(
         default=-1,
         metadata={

paddlenlp/trainer/utils/reshard/common.py

@@ -102,7 +102,7 @@ def convert_opt_name_to_tname(tensor_names, opt_names):
                 opt_to_t[t] = t[: -len(s)]
                 _find = True
                 break
-        assert _find
+        assert _find, t
     return opt_to_t
 
 

paddlenlp/trainer/utils/sharding_io.py

@@ -270,7 +270,7 @@ def get_group_ids(self):
 
 
 class ShardingIO:
-    def __init__(self, args, model, optimizer=None, hcg=None, remap_parameter_name=False):
+    def __init__(self, args, model, optimizer=None, hcg=None, remap_parameter_name=False, is_ema=False):
         self.args = args
         self.model = model
         self.optimizer = optimizer
@@ -282,6 +282,7 @@ def __init__(self, args, model, optimizer=None, hcg=None, remap_parameter_name=F
 
         self.remap_parameter_name = remap_parameter_name
         self.remapper = None
+        self.is_ema = is_ema
 
     def _get_remapper(self, checkpoint):
         if not self.remap_parameter_name:
@@ -395,24 +396,33 @@ def _load_one_state_dict_from_checkpoint(self, resume_from_checkpoint, base_weig
         """
         load state_dict of one shard from_checkpoint, Only load model state dict.
         """
+        if self.is_ema:
+            base_weight_name = base_weight_name.replace("model_state", "ema").replace("pdparams", "pdopt")
         file_path = os.path.join(resume_from_checkpoint, _add_variant(base_weight_name, weight_name_suffix))
         if not os.path.isfile(file_path):
             raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}, no {file_path}")
 
         logger.info(f"Loading model from {resume_from_checkpoint} .")
         # We load the model state dict on the CPU to avoid an OOM error.
         state_dict = paddle.load(file_path, return_numpy=True)
+        if self.is_ema:
+            state_dict.pop("master_weights", None)
         state_dict = self._remap_parameter_name(resume_from_checkpoint, state_dict, is_opt=False)
         return state_dict
 
     def _load_optimizer_state_of_one_shard(self, checkpoint, base_opt_name, optimizer_name_suffix, group_getter=None):
+        if self.is_ema:
+            base_opt_name = base_opt_name.replace("optimizer", "ema")
         optimizer_name = _add_variant(base_opt_name, optimizer_name_suffix)
         path = os.path.join(checkpoint, optimizer_name)
         logger.info(f"load optimizer state from {path}")
         if os.path.isfile(path):
+            opt_state = paddlenlp_load(path, map_location="cpu")
+            if self.is_ema:
+                opt_state = {"master_weights": opt_state.get("master_weights", {})}
             return self._remap_parameter_name(
                 checkpoint,
-                self._modify_ckpt_for_compatibility(paddlenlp_load(path, map_location="cpu")),
+                self._modify_ckpt_for_compatibility(opt_state),
                 is_opt=True,
             )
         logger.info(f"{path} not exists")