[Unified Checkpoint] Support deepep (#2623)

Author: DesmonDay

paddleformers/trainer/training_args.py

@@ -2231,6 +2231,17 @@ def expert_parallel_rank(self):
         else:
             return 0
 
+    @property
+    def moe_sharding_parallel_rank(self):
+        if self.use_hybrid_parallel:
+            hcg = fleet.get_hybrid_communicate_group()
+            if hasattr(hcg, "get_moe_sharding_parallel_group"):
+                return max(hcg.get_moe_sharding_parallel_group().rank, 0)
+            else:
+                return 0
+        else:
+            return 0
+
     def _format_name(self, prefix, rank, degree):
         size = 2
         return f"{prefix}{rank:0>{size}d}"
@@ -2390,7 +2401,9 @@ def should_save_model_state(self):
                 return True
             elif self.use_hybrid_parallel:
                 # save on dataset rank 0
-                return self.sharding_parallel_rank == 0 and (self.data_parallel_rank == 0 or self.use_expert_parallel)
+                return (
+                    self.sharding_parallel_rank == 0 and (self.data_parallel_rank == 0 or self.use_expert_parallel)
+                ) or (self.expert_parallel_degree > 1 and self.moe_sharding_parallel_rank == 0)
             else:
                 return self.process_index == 0 or self.use_expert_parallel
 

paddleformers/trainer/unified_checkpoint/unified_checkpoint.py

@@ -232,10 +232,21 @@ def save_non_merge_optimizer(
             for key in list(master_weights.keys()):
                 master_weights[static2struct_name_mappings[key]] = master_weights.pop(key)
 
-        if self.args.use_expert_parallel:
-            model_state_dict = get_expected_state_dict(model)
-            filter_sync_parameters(model_state_dict, optim_state_dict, is_model_weight=False)
-            filter_sync_parameters(model_state_dict, master_weights, is_model_weight=False)
+        model_state_dict = get_expected_state_dict(model)
+        filter_sync_parameters(
+            model_state_dict,
+            optim_state_dict,
+            is_model_weight=False,
+            use_expert_parallel=self.args.use_expert_parallel,
+            expert_parallel_degree=self.args.expert_parallel_degree,
+        )
+        filter_sync_parameters(
+            model_state_dict,
+            master_weights,
+            is_model_weight=False,
+            use_expert_parallel=self.args.use_expert_parallel,
+            expert_parallel_degree=self.args.expert_parallel_degree,
+        )
 
         optimizer_name = _add_variant(SAFE_OPTIMIZER_NAME, self.args.optimizer_name_suffix)
         master_weights_name = _add_variant(SAFE_MASTER_WEIGHTS_NAME, self.args.optimizer_name_suffix)
@@ -607,9 +618,12 @@ def unified_checkpoint_into_shards(
 
     config_to_save = copy.deepcopy(model_to_save.config)
 
-    if args.use_expert_parallel:
-        # ignore saving `no_sync=False` tensors when using expert_parallel under dp_rank > 0.
-        filter_sync_parameters(state_dict, is_model_weight=True)
+    filter_sync_parameters(
+        state_dict,
+        is_model_weight=True,
+        use_expert_parallel=args.use_expert_parallel,
+        expert_parallel_degree=args.expert_parallel_degree,
+    )
 
     if config_to_save.tensor_parallel_degree > 1:
         if isinstance(model_to_save, LoRAModel) or isinstance(model_to_save, PrefixModelForCausalLM):
@@ -639,20 +653,24 @@ def unified_checkpoint_into_shards(
 
     shard_file = get_sharded_file_name(args, weights_name)
     # renumerize shard_file name for expert_parallel.
-    if args.use_expert_parallel:
+    if args.use_expert_parallel and args.expert_parallel_degree <= 1:
         shard_file = rename_shard_file(args, shard_file, weights_name)
 
     for key, weight in state_dict.items():
         index_weight_file[key] = shard_file
         total_size += weight.numel().item() * dtype_byte_size(weight.dtype)
 
     index_file_list, total_size_list = gather_sharded_object(
-        index_weight_file, total_size, use_expert_parallel=args.use_expert_parallel
+        index_weight_file,
+        total_size,
+        use_expert_parallel=args.use_expert_parallel,
+        expert_parallel_degree=args.expert_parallel_degree,
     )
     sharded_index = get_sharded_index(
         index_file_list,
         total_size_list,
     )
+
     if sharded_index is not None:
         if isinstance(model_to_save, LoRAModel):
             sharded_index["type"] = "lora"
@@ -724,8 +742,13 @@ def unified_optimizer_into_shards(
     tp_group = fleet.get_hybrid_communicate_group().get_model_parallel_group()
     tp_size = tp_group.nranks
 
-    if args.use_expert_parallel:
-        filter_sync_parameters(model_state_dict, state_dict, is_model_weight=False)
+    filter_sync_parameters(
+        model_state_dict,
+        state_dict,
+        is_model_weight=False,
+        use_expert_parallel=args.use_expert_parallel,
+        expert_parallel_degree=args.expert_parallel_degree,
+    )
 
     if tp_size > 1:
         # get tp_actions

paddleformers/trainer/unified_checkpoint/utils.py

@@ -493,7 +493,8 @@ def filter_params(model_to_save, state_dict, args, is_optimizer=False):
             weight_key = k.split("/")[0]
             model_v = model_state_dict[weight_key] if is_optimizer else v
             mp_moe = getattr(model_v, "mp_moe", False)
-            if not mp_moe:
+            no_sync = getattr(model_v, "no_sync", False)
+            if not mp_moe or no_sync:
                 if not quant or not is_optimizer:
                     if hasattr(model_v, "is_distributed") and model_v.is_distributed:
                         tensor_bytes_dict[k] = v.numel().item() * tp_size * dtype_byte_size(v.dtype)
@@ -555,6 +556,9 @@ def filter_params(model_to_save, state_dict, args, is_optimizer=False):
         mp_moe = getattr(model_v, "mp_moe", False)
         if mp_moe:
             filter_tensor_list[tp_rank].append(k)
+        no_sync = getattr(model_v, "no_sync", False)
+        if no_sync and k not in filter_tensor_list[tp_rank]:
+            filter_tensor_list[tp_rank].append(k)
 
     final_filter_tensor_list = []
     dist.all_gather_object(final_filter_tensor_list, filter_tensor_list[tp_rank], group=tp_group)
@@ -568,14 +572,20 @@ def get_sharded_file_name(args, file_name, is_optimizer=False):
     """
     if not is_optimizer:
         sd_degree = args.sharding_parallel_degree if args.sharding_parallel_degree > 1 else 1
-        size = sd_degree if args.use_expert_parallel else args.dataset_world_size
+        if args.use_expert_parallel:
+            if args.expert_parallel_degree > 1:
+                size = dist.get_world_size() // args.moe_sharding_parallel_degree
+            else:
+                size = args.world_size // sd_degree
+        else:
+            size = args.world_size // args.dataset_world_size
         shard_file = file_name.replace(
             ".pdparams",
-            f"-{args.logical_process_index + 1:05d}-of-{args.world_size//size:05d}.pdparams",
+            f"-{args.logical_process_index + 1:05d}-of-{size:05d}.pdparams",
         )
         shard_file = shard_file.replace(
             ".safetensors",
-            f"-{args.logical_process_index + 1:05d}-of-{args.world_size//size:05d}.safetensors",
+            f"-{args.logical_process_index + 1:05d}-of-{size:05d}.safetensors",
         )
     else:
         hcg = fleet.get_hybrid_communicate_group()
@@ -617,7 +627,9 @@ def get_sharded_index(
     return None
 
 
-def gather_sharded_object(index_file, total_size, is_optimizer=False, use_expert_parallel=False):
+def gather_sharded_object(
+    index_file, total_size, is_optimizer=False, use_expert_parallel=False, expert_parallel_degree=1
+):
     """
     All gather sharded files list across different groups.
     """
@@ -654,7 +666,7 @@ def gather_sharded_object(index_file, total_size, is_optimizer=False, use_expert
         index_file_list = [index_file]
         total_size_list = [total_size]
 
-    if use_expert_parallel:
+    if use_expert_parallel and expert_parallel_degree <= 1:
         data_group = hcg.get_data_parallel_group()
         if data_group.nranks > 1:
             data_index_file_list = []
@@ -664,7 +676,7 @@ def gather_sharded_object(index_file, total_size, is_optimizer=False, use_expert
             index_file_list = flatten_list(data_index_file_list)
             total_size_list = flatten_list(data_total_size_list)
 
-    if is_optimizer:
+    if is_optimizer or expert_parallel_degree > 1:
         sharding_group = hcg.get_sharding_parallel_group()
         if sharding_group.nranks > 1:
             sharding_index_file_list = []
@@ -781,29 +793,48 @@ def save_config(model_to_save):
         model_to_save.generation_config.save_pretrained(save_directory)
 
 
-def filter_sync_parameters(model_state_dict, optim_state_dict=None, master_weights=None, is_model_weight=True):
+def filter_sync_parameters(
+    model_state_dict,
+    optim_state_dict=None,
+    master_weights=None,
+    is_model_weight=True,
+    use_expert_parallel=False,
+    expert_parallel_degree=1,
+):
     """Filter sync parameters under expert parallel mode."""
 
     hcg = fleet.get_hybrid_communicate_group()
     dp_group = hcg.get_data_parallel_group()
+    sharding_group = hcg.get_sharding_parallel_group()
     dp_rank = dp_group.rank if dp_group.nranks > 1 else 0
+    sharding_rank = sharding_group.rank if sharding_group.nranks > 1 else 0
+    if expert_parallel_degree > 1:
+        ep_group = hcg.get_expert_parallel_group()
+        ep_rank = ep_group.rank if ep_group.nranks > 1 else 0
+    logger.info("Filter sync parameters under expert parallel mode.")
 
     if is_model_weight:
         for key in list(model_state_dict.keys()):
-            if dp_rank > 0 and not getattr(model_state_dict[key], "no_sync", False):
-                model_state_dict.pop(key)
+            if use_expert_parallel:
+                if expert_parallel_degree > 1:
+                    if ep_rank > 0 and sharding_rank > 0 and not getattr(model_state_dict[key], "no_sync", False):
+                        model_state_dict.pop(key)
+                else:
+                    if dp_rank > 0 and not getattr(model_state_dict[key], "no_sync", False):
+                        model_state_dict.pop(key)
     else:
-        no_sync_kname = []
-        for k, v in model_state_dict.items():
-            if getattr(v, "no_sync", False):
-                no_sync_kname.append(k)
-
-        for key in list(optim_state_dict.keys()):
-            model_key = key.split("/")[0]
-            if dp_rank > 0 and model_key not in no_sync_kname:
-                optim_state_dict.pop(key)
-
-        if master_weights is not None:
-            for key in list(master_weights.keys()):
-                if dp_rank > 0 and key not in no_sync_kname:
-                    master_weights.pop(key)
+        if use_expert_parallel and expert_parallel_degree == 1:
+            no_sync_kname = []
+            for k, v in model_state_dict.items():
+                if getattr(v, "no_sync", False):
+                    no_sync_kname.append(k)
+
+            for key in list(optim_state_dict.keys()):
+                model_key = key.split("/")[0]
+                if dp_rank > 0 and model_key not in no_sync_kname:
+                    optim_state_dict.pop(key)
+
+            if master_weights is not None:
+                for key in list(master_weights.keys()):
+                    if dp_rank > 0 and key not in no_sync_kname:
+                        master_weights.pop(key)

paddleformers/transformers/moe_layer.py

@@ -353,6 +353,7 @@ def __init__(self, config, moe_num_experts, expert_class, expert_kwargs, gate, m
             self.num_local_experts, self.moe_router_topk, self.moe_num_experts, moe_group
         )
         self.expert_parallel_degree = 1 if self.ep_size < 0 else self.ep_size
+        self.is_dummy_moe = False if self.expert_parallel_degree > 1 else True
         self.moe_num_experts_per_device = self._parse_moe_expert_parallel(
             self.moe_num_experts, self.expert_parallel_degree
         )
@@ -363,6 +364,14 @@ def __init__(self, config, moe_num_experts, expert_class, expert_kwargs, gate, m
             else:
                 self.experts.append(None)
         self.gate = gate
+        self._post_init()
+
+    def _post_init(self):
+        for k in self.experts:
+            if k is not None:
+                for p in k.parameters():
+                    p.expert = not self.is_dummy_moe
+                    p.no_sync = not self.is_dummy_moe
 
     def expert_forward(self, dispatched_input, tokens_per_expert):
         outputs = []