Cherry-pick hybrid expert parallel sharding_metas (#2447)

Author: pkuzyc

paddleformers/trainer/trainer.py

@@ -377,7 +377,13 @@ def __init__(
         self.optimizer_grouped_parameters = None
         self.sharding_io = None
         if self.args.should_save_sharding_stage1_model or self.args.should_load_sharding_stage1_model:
-            self.sharding_io = ShardingIO(self.args, self.model, self.optimizer)
+            self.sharding_io = ShardingIO(
+                self.args,
+                self.model,
+                self.optimizer,
+                remap_parameter_name=self.args.load_sharded_model_remap_parameter_name,
+            )
+
         if self.args.unified_checkpoint:
             self.unified_checkpoint_handler = UnifiedCheckpointHandler(self.args)
 

paddleformers/trainer/training_args.py

@@ -627,6 +627,11 @@ class TrainingArguments:
         },
     )
 
+    load_sharded_model_remap_parameter_name: bool = field(
+        default=False,
+        metadata={"help": "Whether to remap parameter name when load_sharded_model = true."},
+    )
+
     tensor_parallel_degree: int = field(
         default=-1,
         metadata={

paddleformers/trainer/utils/reshard/__init__.py

@@ -18,6 +18,7 @@
     SHARDING_STRATEGY_V2,
     NodeModelState,
     all_gather_state_dict,
+    convert_opt_name_to_tname,
     get_moe_sharding_group,
     get_param_sharding_group,
     get_sharding_strategy,

paddleformers/trainer/utils/reshard/common.py

@@ -22,6 +22,8 @@
 )
 from paddle.distributed.fleet.utils.log_util import logger
 
+from paddleformers.utils.tools import get_env_device
+
 try:
     from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
         DygraphShardingOptimizerV2,
@@ -61,6 +63,49 @@ def get_sharding_strategy(optimizer):
     return SHARDING_STRATEGY_V1
 
 
+def convert_opt_name_to_tname(tensor_names, opt_names):
+    tensor_names = set(tensor_names)
+    all_names = []
+    all_names.extend(list(tensor_names))
+    all_names.extend(opt_names)
+    all_names.sort()
+    pre_t_name = ""
+    suffix = [
+        "_fp32_master_0_beta1_pow_acc_0",
+        "_fp32_master_0_beta2_pow_acc_0",
+        "_fp32_master_0_moment1_0",
+        "_fp32_master_0_moment2_0",
+        "_beta1_pow_acc_0",
+        "_beta2_pow_acc_0",
+        "_moment1_0",
+        "_moment2_0",
+    ]
+    opt_to_t = {}
+    for n in all_names:
+        if n in tensor_names:
+            # we get a param
+            pre_t_name = n
+        else:
+            assert pre_t_name
+            opt_to_t[n] = pre_t_name
+
+    for t in opt_names:
+        _find = False
+        for s in suffix:
+            if get_env_device() == "xpu" and t.endswith(s + ".SCALE_VALUE"):
+                # NOTE: for xpu adamw, all optimizer state will have an extra attribute end with SCALE_VALUE.
+                # This extra attribute won't be used, just skip it.
+                _find = True
+                break
+            if t.endswith(s):
+                logger.info(f"{t}-{t[:-len(s)]}--{t[:-len(s)] in tensor_names}")
+                opt_to_t[t] = t[: -len(s)]
+                _find = True
+                break
+        assert _find
+    return opt_to_t
+
+
 class NodeModelState:
     def __init__(self, group):
         self._model_weights = OrderedDict()
@@ -259,43 +304,6 @@ def pack_keys(self, structure_name_mapping=None):
         change the key of master weights dict from param_name to (structure_name, param_name)
         """
         # pack key for pp convert
-        def _opt_name_to_tname(tensor_names, opt_names):
-            tensor_names = set(tensor_names)
-            all_names = []
-            all_names.extend(list(tensor_names))
-            all_names.extend(opt_names)
-            all_names.sort()
-            pre_t_name = ""
-            suffix = [
-                "_fp32_master_0_beta1_pow_acc_0",
-                "_fp32_master_0_beta2_pow_acc_0",
-                "_fp32_master_0_moment1_0",
-                "_fp32_master_0_moment2_0",
-                "_beta1_pow_acc_0",
-                "_beta2_pow_acc_0",
-                "_moment1_0",
-                "_moment2_0",
-            ]
-            opt_to_t = {}
-            for n in all_names:
-                if n in tensor_names:
-                    # we get a param
-                    pre_t_name = n
-                else:
-                    assert pre_t_name
-                    opt_to_t[n] = pre_t_name
-
-            for t in opt_names:
-                _find = False
-                for s in suffix:
-                    if t.endswith(s):
-                        logger.info(f"{t}-{t[:-len(s)]}--{t[:-len(s)] in tensor_names}")
-                        opt_to_t[t] = t[: -len(s)]
-                        _find = True
-                        break
-                assert _find
-            return opt_to_t
-
         if structure_name_mapping is not None:
             tname_to_structure_name = {v: k for (k, v) in structure_name_mapping.items()}
         else:
@@ -304,7 +312,7 @@ def _opt_name_to_tname(tensor_names, opt_names):
 
         tensor_names = list(tname_to_structure_name.keys())
         opt_names = list(self._opt_state.keys())
-        opt_name_to_tname = _opt_name_to_tname(tensor_names, opt_names)
+        opt_name_to_tname = convert_opt_name_to_tname(tensor_names, opt_names)
 
         # model state
         model_weights_tmp = OrderedDict()