support tpdp-ep sharding reshard

Author: sneaxiy

paddlenlp/trainer/training_args.py

@@ -1230,6 +1230,9 @@ def __post_init__(self):
 
             if expert_parallel_degree > 1:
                 moe_sharding_parallel_degree = world_size // (pipeline_parallel_degree * expert_parallel_degree)
+                assert (
+                    self.expert_tensor_parallel_degree <= 1
+                ), "expert_tensor_parallel_degree > 1 is not supported when expert_parallel_degree > 1"
             else:
                 moe_sharding_parallel_degree = 1
             moe_sharding_parallel_degree = max(moe_sharding_parallel_degree, 1)
@@ -2164,6 +2167,17 @@ def pipeline_parallel_rank(self):
         else:
             return 0
 
+    @property
+    def expert_parallel_rank(self):
+        if self.use_hybrid_parallel:
+            hcg = fleet.get_hybrid_communicate_group()
+            if hasattr(hcg, "get_expert_parallel_rank"):
+                return max(hcg.get_expert_parallel_rank(), 0)
+            else:
+                return 0
+        else:
+            return 0
+
     def _format_name(self, prefix, rank, degree):
         size = 2
         return f"{prefix}{rank:0>{size}d}"
@@ -2178,7 +2192,7 @@ def optimizer_name_suffix(self):
                 name.append(self._format_name("pp", self.pipeline_parallel_rank, self.pipeline_parallel_degree))
             if self.sharding_parallel_degree > 1:
                 name.append(self._format_name("shard", self.sharding_parallel_rank, self.sharding_parallel_degree))
-            if self.use_expert_parallel:
+            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
                 name.append(self._format_name("moe", self.data_parallel_rank, self.data_parallel_degree))
             return "_".join(name)
         else:
@@ -2194,7 +2208,7 @@ def weight_name_suffix(self):
                 name.append(self._format_name("tp", self.tensor_parallel_rank, self.tensor_parallel_degree))
             if self.pipeline_parallel_degree > 1:
                 name.append(self._format_name("pp", self.pipeline_parallel_rank, self.pipeline_parallel_degree))
-            if self.use_expert_parallel:
+            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
                 name.append(self._format_name("moe", self.data_parallel_rank, self.data_parallel_degree))
             return "_".join(name)
 
@@ -2220,7 +2234,7 @@ def sharded_name_suffix(self, shard_id=None, pp_id=None, moe_id=None, sharding_p
                     shard_id = self.sharding_parallel_rank
                 assert isinstance(shard_id, int)
                 name.append(self._format_name("shard", shard_id, sharding_parallel_degree))
-            if self.use_expert_parallel:
+            if self.use_expert_parallel and self.expert_parallel_degree <= 1:
                 if moe_id is None:
                     moe_id = self.data_parallel_rank
                 assert isinstance(moe_id, int)

paddlenlp/trainer/utils/reshard/__init__.py

@@ -19,6 +19,13 @@
     NodeModelState,
     all_gather_state_dict,
     convert_opt_name_to_tname,
+    get_moe_sharding_group,
+    get_param_sharding_group,
     get_sharding_strategy,
     is_sharding_opt,
+    merge_model_state,
+    merge_opt_state,
+    split_model_state,
+    split_opt_state,
+    split_structure_name_mapping,
 )

paddlenlp/trainer/utils/reshard/common.py

@@ -16,6 +16,7 @@
 
 import numpy as np
 import paddle
+import paddle.distributed.fleet as fleet
 from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
     DygraphShardingOptimizer,
 )
@@ -106,17 +107,16 @@ def convert_opt_name_to_tname(tensor_names, opt_names):
 
 
 class NodeModelState:
-    def __init__(self, mp_rank=None, sharding_rank=None, pp_rank=None):
+    def __init__(self, group):
         self._model_weights = OrderedDict()
         self._opt_state = OrderedDict()
         self._master_weights = OrderedDict()
         self._lr_scheduler = None
-        self.set_node_rank(mp_rank, sharding_rank, pp_rank)
+        self._group = group
 
-    def set_node_rank(self, mp_rank, sharding_rank, pp_rank):
-        self._mp_rank = mp_rank
-        self._sharding_rank = sharding_rank
-        self._pp_rank = pp_rank
+    @property
+    def group(self):
+        return self._group
 
     def _add_kv(self, d, k, v):
         assert k not in d
@@ -407,12 +407,13 @@ def split_state(self, split_func):
 
         return node_model_states
 
-    def even_distribute(self, group):
+    def even_distribute(self):
         """
         distribute the node state evenly among all workers in group， and make sure
         in the dicts of (key, rank)=>tensor, items keys of the same key but different rank are distributed to the
         same worker
         """
+        group = self.group
         # sharding degree == 1
         if group is None or group.nranks < 2:
             return self
@@ -446,7 +447,7 @@ def distribute(state_dict):
             def filter_func(key):
                 assert key[0] in key_to_rank, key
                 dst_rank = key_to_rank[key[0]]
-                return dst_rank == group.rank
+                return dst_rank == max(group.rank, 0)
 
             return _all_gather_state_dict(state_dict, filter_func, group)
 
@@ -455,10 +456,11 @@ def filter_func(key):
         self._master_weights = distribute(self._master_weights)
         return self
 
-    def reshard(self, group, filter_func):
+    def reshard(self, filter_func):
         """
         reshard according to the passed in filter_func
         """
+        group = self.group
         self._model_weights = _all_gather_state_dict(self._model_weights, filter_func, group)
         self._opt_state = _all_gather_state_dict(self._opt_state, filter_func, group)
         self._master_weights = _all_gather_state_dict(self._master_weights, filter_func, group)
@@ -511,6 +513,7 @@ def merge(state, l):
         return self
 
     def merge_from(self, other, rank=None):
+        assert other.group is self.group
         self.add_weights(other.model_weights, rank)
         self.add_opts(other.opt_state, rank)
         self.add_master_weights(other.master_weights, rank)
@@ -528,6 +531,68 @@ def get_opt_state_dict(self):
         return opt_state_dict
 
 
+def split_model_state(model_state, group_getter):
+    res = OrderedDict()
+    for k, v in model_state.items():
+        group = group_getter.get_group(k)
+        if group.id not in res:
+            res[group.id] = OrderedDict()
+        res[group.id][k] = v
+    return res
+
+
+def merge_model_state(model_state_map):
+    res = OrderedDict()
+    for gid, model_state in model_state_map.items():
+        res.update(model_state)
+    return res
+
+
+def split_opt_state(opt_state, group_getter):
+    res = OrderedDict()
+    lr_scheduler = opt_state.get("LR_Scheduler", None)
+    for k, v in opt_state.items():
+        if k == "LR_Scheduler":
+            continue
+        elif k == "master_weights":
+            for kk, vv in v.items():
+                group = group_getter.get_group(kk)
+                if group.id not in res:
+                    res[group.id] = {"master_weights": OrderedDict(), "LR_Scheduler": lr_scheduler}
+                res[group.id]["master_weights"][kk] = vv
+        else:
+            assert isinstance(v, paddle.Tensor), type(v)
+            group = group_getter.get_group(k)
+            if group.id not in res:
+                res[group.id] = {"master_weights": OrderedDict(), "LR_Scheduler": lr_scheduler}
+            res[group.id][k] = v
+    return res
+
+
+def merge_opt_state(opt_state_map):
+    res = {"LR_Scheduler": None, "master_weights": OrderedDict()}
+    for gid, opt_state in opt_state_map.items():
+        for k, v in opt_state.items():
+            if k == "LR_Scheduler":
+                if v is not None:
+                    res["LR_Scheduler"] = v
+            elif k == "master_weights":
+                res["master_weights"].update(v)
+            else:
+                res[k] = v
+    return res
+
+
+def split_structure_name_mapping(structure_name_mapping, group_getter):
+    res = OrderedDict()
+    for k, v in structure_name_mapping.items():
+        group = group_getter.get_group(k)
+        if group.id not in res:
+            res[group.id] = OrderedDict()
+        res[group.id][k] = v
+    return res
+
+
 def all_gather_simple_object(obj, group):
     res = []
     if group.nranks < 2:
@@ -570,7 +635,7 @@ def map_func(weight):
             del state_dict[k]
         else:
             tensor = paddle.to_tensor(np.empty(shape, dtype))
-        logger.info(f"broadcast {k} from {rank}")
+        logger.info(f"broadcast {k} from {rank}, group {group}")
         # broadcast the tensor
         if group.nranks > 1:
             paddle.distributed.broadcast(
@@ -595,3 +660,29 @@ def _all_gather_state_dict(state_dict, filter_func, group):
     for (k, v) in tmp_state_dict.items():
         state_dict[k] = v
     return state_dict
+
+
+def get_moe_sharding_group(hcg=None):
+    if hcg is None:
+        hcg = fleet.get_hybrid_communicate_group()
+    if hasattr(hcg, "get_moe_sharding_parallel_group"):
+        return hcg.get_moe_sharding_parallel_group()
+    else:
+        return None
+
+
+def get_param_sharding_group(param, hcg=None):
+    if hcg is None:
+        hcg = fleet.get_hybrid_communicate_group()
+    default_group = hcg.get_sharding_parallel_group()
+    ep_sharding_group = get_moe_sharding_group(hcg)
+
+    if not hasattr(param, "color"):
+        return default_group
+    color = getattr(param, "color")
+    if isinstance(color, dict):
+        group = color.get("group", default_group)
+        assert group is default_group or group is ep_sharding_group, f"unsupported group: {group}"
+        return group
+    else:
+        return default_group

paddlenlp/trainer/utils/reshard/sharding_v1.py

@@ -20,9 +20,8 @@
 from .common import is_sharding_opt
 
 
-def shard(node_model_state, model, optimizer, hcg):
-    group = hcg.get_sharding_parallel_group()
-    cur_rank = group.rank
+def shard(node_model_state, model, optimizer):
+    cur_rank = max(node_model_state.group.rank, 0)
     unwrapped_optimizer = unwrap_optimizer(optimizer, DygraphShardingOptimizer)
     if unwrapped_optimizer is not None:
         optimizer = unwrapped_optimizer
@@ -40,10 +39,10 @@ def filter_func(key):
         assert not is_sharding_opt(optimizer)
         filter_func = lambda key: True
 
-    node_model_state.reshard(group, filter_func)
+    node_model_state.reshard(filter_func)
     return node_model_state
 
 
-def restore(node_model_state, model, optimizer, hcg):
+def restore(node_model_state, model, optimizer):
     node_model_state.drop_rank()
     return node_model_state

paddlenlp/trainer/utils/reshard/sharding_v2.py

@@ -34,13 +34,16 @@
 
 from paddle.distributed.communication.reduce import ReduceOp
 
+from .common import get_moe_sharding_group
 
-def shard(node_model_state, model, optimizer, hcg):
+
+def shard(node_model_state, model, optimizer):
     assert DygraphShardingOptimizerV2 is not None
-    group = hcg.get_sharding_parallel_group()
-    cur_rank = group.rank
     split_infos = collect_split_info(optimizer, model)
 
+    group = node_model_state.group
+    cur_rank = max(group.rank, 0)
+
     def split_func(k, v):
         param_name = k[1]
         opt_name = k[-1]
@@ -87,15 +90,14 @@ def filter_func(k):
         return rank == cur_rank
 
     # reshard
-    node_model_state.reshard(group, filter_func)
+    node_model_state.reshard(filter_func)
     node_model_state.drop_rank()
     return node_model_state
 
 
-def restore(node_model_state, model, optimizer, hcg):
-    group = hcg.get_sharding_parallel_group()
+def restore(node_model_state, model, optimizer):
     # evenly distribute param
-    node_model_state.even_distribute(group)
+    node_model_state.even_distribute()
     param_shapes = {k: v.shape for (k, v) in model.state_dict().items()}
 
     def merge_func(k, v):
@@ -175,7 +177,7 @@ def gather_infos(comm_buffer):
         for comm_buffer in optimizer._comm_buffer_list:
             gather_infos(comm_buffer)
 
-    assert len(split_infos)
+    assert len(split_infos) > 0
     return split_infos
 
 
@@ -211,11 +213,16 @@ def get_matched_length(name):
     if need_allgather:
         if hcg is None:
             hcg = fleet.get_hybrid_communicate_group()
-        group = hcg.get_sharding_parallel_group()
-        if group is not None and group.nranks > 1:
-            x = paddle.to_tensor([is_matched], dtype=paddle.int32)
-            paddle.distributed.stream.all_reduce(x, op=ReduceOp.MIN, group=group, sync_op=True, use_calc_stream=True)
-            global_is_matched = int(x.numpy()[0])
+        sharding_group = hcg.get_sharding_parallel_group()
+        moe_sharding_group = get_moe_sharding_group(hcg)
+        for group in [sharding_group, moe_sharding_group]:
+            if group is not None and group.nranks > 1:
+                x = paddle.to_tensor([is_matched], dtype=paddle.int32)
+                paddle.distributed.stream.all_reduce(
+                    x, op=ReduceOp.MIN, group=group, sync_op=True, use_calc_stream=True
+                )
+                is_matched = int(x.numpy()[0])
+        global_is_matched = is_matched
     else:
         global_is_matched = is_matched