Cherry-pick some PRs from PaddleNLP (#2821)

Author: sneaxiy

paddleformers/trainer/trainer.py

@@ -394,6 +394,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:
             self.unified_checkpoint_handler = UnifiedCheckpointHandler(self.args)
@@ -836,9 +837,16 @@ def create_zcc_manager(self, unwrapped_model, resume_from_checkpoint=None):
         if resume_from_checkpoint is not None:
             path = _add_variant(PADDLE_OPTIMIZER_NAME, self.args.optimizer_name_suffix)
             path = os.path.join(resume_from_checkpoint, path).replace("optimizer", "ema")
+            if self.args.zcc_save_ema_coef is not None and self.sharding_io is not None:
+                success, err_msg = self.sharding_io.check_same_strategy(resume_from_checkpoint)
+            else:
+                success, err_msg = True, None
             if os.path.exists(path):
-                logger.info(f"ZCC EMA load from {path}")
-                self.zcc_manager.set_ema_state_dict(path)
+                if success:
+                    logger.info(f"ZCC EMA load from {path}")
+                    self.zcc_manager.set_ema_state_dict(path)
+                else:
+                    logger.info(f"ZCC EMA does not load {path} because {err_msg}")
             else:
                 logger.info(f"ZCC EMA state dict not found, in: {path}")
 

paddleformers/trainer/training_args.py

@@ -634,6 +634,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={
@@ -2504,9 +2509,7 @@ def should_save_sharding_stage1_model(self):
     def should_load_sharding_stage1_model(self):
         if self.enable_auto_parallel:
             return False
-        return (
-            ShardingOption.SHARD_OP in self.sharding and self.sharding_parallel_degree > 1 and self.load_sharded_model
-        )
+        return self.load_sharded_model
 
     @property
     def should_load_dataset(self):

paddleformers/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
 
 
@@ -609,12 +609,13 @@ def map_func(weight):
             weight = weight.numpy()
         return weight
 
+    group_rank = max(group.rank, 0)
     state_dict = {k: map_func(v) for (k, v) in state_dict.items()}
 
     meta_dict = {}
     for (k, v) in state_dict.items():
         # src rank
-        meta_dict[k] = (v.dtype, v.shape, group.rank)
+        meta_dict[k] = (v.dtype, v.shape, group_rank)
 
     meta_dict_list = all_gather_simple_object(meta_dict, group)
 
@@ -628,20 +629,21 @@ def map_func(weight):
     meta_list = sorted(meta_list, key=lambda x: x[0])
     for (k, meta) in meta_list:
         dtype, shape, rank = meta
-        if rank == group.rank:
+        if rank == group_rank:
             assert k in state_dict
             tensor = paddle.to_tensor(state_dict[k])
             del state_dict[k]
         else:
             tensor = paddle.to_tensor(np.empty(shape, dtype))
         logger.info(f"broadcast {k} from {rank}, group {group}")
         # broadcast the tensor
-        paddle.distributed.broadcast(
-            tensor,
-            src=group.ranks[rank],
-            group=group,
-            sync_op=True,
-        )
+        if group.nranks > 1:
+            paddle.distributed.broadcast(
+                tensor,
+                src=group.ranks[rank],
+                group=group,
+                sync_op=True,
+            )
         if filter_func(k):
             res[k] = tensor.cpu()
         del tensor

paddleformers/trainer/utils/reshard/sharding_v1.py

@@ -17,20 +17,27 @@
 )
 
 from ....transformers.model_utils import unwrap_optimizer
+from .common import is_sharding_opt
 
 
 def shard(node_model_state, model, optimizer):
     cur_rank = max(node_model_state.group.rank, 0)
-    optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
-    assert optimizer is not None
-    param2rank = optimizer._param2rank
-
-    def filter_func(key):
-        names = key
-        param_name = names[1]
-        assert param_name in param2rank
-        dst_rank = param2rank[param_name]
-        return dst_rank == cur_rank
+    unwrapped_optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
+    if unwrapped_optimizer is not None:
+        optimizer = unwrapped_optimizer
+        assert not is_sharding_opt(optimizer)
+        param2rank = optimizer._param2rank
+
+        def filter_func(key):
+            names = key
+            param_name = names[1]
+            assert param_name in param2rank
+            dst_rank = param2rank[param_name]
+            return dst_rank == cur_rank
+
+    else:
+        assert not is_sharding_opt(optimizer)
+        filter_func = lambda key: True
 
     node_model_state.reshard(filter_func)
     return node_model_state

paddleformers/trainer/utils/sharding_io.py

@@ -269,7 +269,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
@@ -281,6 +281,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:
@@ -351,7 +352,9 @@ def load_model_slices():
                 structure_name_map = split_structure_name_mapping(structure_name_map, group_getter)
                 for i in range(self.args.sharding_parallel_rank, sharding_degree, cur_sharding_degree):
                     tmp = self._load_one_state_dict_from_checkpoint(
-                        checkpoint, base_weight_name, self.args.sharded_name_suffix(i, j)
+                        checkpoint,
+                        base_weight_name,
+                        self.args.sharded_name_suffix(i, j, sharding_parallel_degree=sharding_degree),
                     )
                     tmp = split_model_state(tmp, group_getter)
                     for gid in gids:
@@ -399,24 +402,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 = paddleformers_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(paddleformers_load(path, map_location="cpu")),
+                self._modify_ckpt_for_compatibility(opt_state),
                 is_opt=True,
             )
         logger.info(f"{path} not exists")
@@ -449,9 +461,12 @@ def _need_reshard(self, checkpoint):
         if sharding_strategy == SHARDING_STRATEGY_V1:
             param2rank = sharding_meta["param2rank"]
             optimizer = unwrap_optimizer(self.optimizer, DygraphShardingOptimizer)
-            assert optimizer
-            if len(param2rank) == 0:
-                logger.warning("The param2rank is empty. Force reshard would be performed.")
+            if self.args.sharding_parallel_degree > 1:
+                assert optimizer is not None
+            else:
+                assert optimizer is None
+            if len(param2rank) == 0 or optimizer is None:
+                logger.warning("The param2rank is empty or sharding degree is 1. Force reshard would be performed.")
                 return True
             assert len(param2rank) == len(optimizer._param2rank)
             for (k, v) in param2rank.items():