support ep for fsdp

Author: zigzagcai

internlm/checkpoint/checkpoint_manager.py

@@ -582,7 +582,7 @@ def try_resume_training(self, train_state: TrainState, current_time=""):
                     f"tp={gpc.get_local_rank(ParallelMode.TENSOR)},pp={gpc.get_local_rank(ParallelMode.PIPELINE)},"
                     f"dp={gpc.get_local_rank(ParallelMode.DATA)}==========="
                 )
-        elif is_using_fsdp() and is_using_hf() and not self.auto_resume:
+        elif is_using_fsdp() and not self.auto_resume:
             pass
         else:
             load_path = self.load_ckpt_info["path"]

internlm/core/fsdp.py

@@ -33,8 +33,10 @@
     FSDP2_SUPPORTED = False
 
 try:
+    import torch.distributed.checkpoint as dcp
     from torch.distributed.checkpoint.state_dict import (
         StateDictOptions,
+        get_model_state_dict,
         set_model_state_dict,
     )
 
@@ -163,8 +165,27 @@ def wrap_FSDP_model(model: Union[nn.Module, nn.ModuleList]):
         )
         fsdp_mode = gpc.config.parallel.fsdp.get("mode", "v1")
         fsdp_init_method = gpc.config.parallel.fsdp.get("init_method", "cuda")
+        
+        if gpc.is_using_parallel_mode(ParallelMode.EXPERT):
+            assert gpc.get_world_size(ParallelMode.EXPERT_DATA) * gpc.get_world_size(ParallelMode.EXPERT) == gpc.get_world_size(ParallelMode.GLOBAL)
 
         if fsdp_mode == "v1":
+            ignored_mod = []
+            if gpc.is_using_parallel_mode(ParallelMode.EXPERT):
+                for layer_id, layer in enumerate(model.model.layers):
+                    if layer_id >= gpc.config.model.first_k_dense_replace:
+                        layer.feed_forward.moe_layer.experts = FSDP(
+                            layer.feed_forward.moe_layer.experts, 
+                            process_group=gpc.get_group(ParallelMode.EXPERT_DATA),
+                            sharding_strategy=ShardingStrategy.FULL_SHARD, 
+                            sync_module_states=fsdp_init_method != "cuda",  # sync model paramters
+                            forward_prefetch=True,
+                            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                            limit_all_gathers=True,
+                            use_orig_params=True,
+                            device_id=None if fsdp_init_method == "cuda" else get_current_device(),  # needed for sync_module_states
+                        )
+                        ignored_mod.append(layer.feed_forward.moe_layer.experts)
             model = FSDP(
                 module=model,
                 process_group=gpc.get_group(ParallelMode.GLOBAL),
@@ -176,6 +197,7 @@ def wrap_FSDP_model(model: Union[nn.Module, nn.ModuleList]):
                 limit_all_gathers=True,
                 use_orig_params=True,
                 device_id=None if fsdp_init_method == "cuda" else get_current_device(),  # needed for sync_module_states
+                ignored_modules=ignored_mod,
             )
             # For FSDP v1, to get ckpt resuming work normally, we do dummy forward.
             # This hack is needed due to FSDP v1 lazy initialization in model construction.
@@ -196,7 +218,7 @@ def wrap_FSDP_model(model: Union[nn.Module, nn.ModuleList]):
         else:
             raise ValueError(f"Unsupported FSDP mode: {fsdp_mode}")
 
-        if is_using_hf() and not gpc.config.ckpt.get("auto_resume", False):
+        if not gpc.config.ckpt.get("auto_resume", False):
             load_ckpt_info = gpc.config.ckpt.load_ckpt_info
             load_ckpt_path = load_ckpt_info.get("path", None)
             load_ckpt_content = load_ckpt_info.get("content", [])
@@ -205,19 +227,25 @@ def wrap_FSDP_model(model: Union[nn.Module, nn.ModuleList]):
                     "model",
                 ), "If auto_resume=False and checkpoint path is given, only model can be loaded"
                 if DCP_SUPPORTED:
-                    hf = gpc.config.hf
-                    mod = LazyObject(hf.mod, hf.mod_cls)
-                    mod = mod.build()
-                    state_dict = mod.from_pretrained(
-                        pretrained_model_name_or_path=load_ckpt_path, use_safetensors=True
-                    ).state_dict()
-                    state_dict = {f"model.{key}": state_dict[key].clone().detach() for key in state_dict}
-                    set_model_state_dict(
-                        model=model, model_state_dict=state_dict, options=StateDictOptions(full_state_dict=True)
-                    )
+                    if is_using_hf():
+                        hf = gpc.config.hf
+                        mod = LazyObject(hf.mod, hf.mod_cls)
+                        mod = mod.build()
+                        state_dict = mod.from_pretrained(
+                            pretrained_model_name_or_path=load_ckpt_path, use_safetensors=True
+                        ).state_dict()
+                        state_dict = {f"model.{key}": state_dict[key].clone().detach() for key in state_dict}
+                        set_model_state_dict(
+                            model=model, model_state_dict=state_dict, options=StateDictOptions(full_state_dict=True)
+                        )
+                    else:
+                        state_dict = get_model_state_dict(model=model)
+                        state_dict = {key: state_dict[key].clone().detach() for key in state_dict}
+                        dcp.load(state_dict=state_dict, checkpoint_id=load_ckpt_path)
+                        set_model_state_dict(model=model, model_state_dict=state_dict)
                     del state_dict
                     internlm_accelerator.empty_cache()
                 else:
                     raise RuntimeError("DCP is not supported in this version of PyTorch.")
 
-    return model
+    return model

internlm/initialize/initialize_optimizer.py

@@ -50,7 +50,7 @@ def split_params_into_different_groups_for_optimizer(
 
     if is_using_fsdp():
         optimizer_mode = ParallelMode.GLOBAL
-        optimizer_mode_expert = ParallelMode.GLOBAL
+        optimizer_mode_expert = ParallelMode.EXPERT_DATA
     else:
         optimizer_mode = ParallelMode.ZERO1
         optimizer_mode_expert = ParallelMode.EXPERT_DATA

internlm/solver/optimizer/fsdp_optimizer.py

@@ -16,7 +16,7 @@
     get_norm,
     release_param_grad,
 )
-from internlm.utils.common import get_tensor_norm, move_norm_to_cuda
+from internlm.utils.common import get_current_device, get_tensor_norm, move_norm_to_cuda
 from internlm.utils.config import Config
 from internlm.utils.logger import get_logger
 
@@ -37,6 +37,7 @@
 def compute_norm(
     gradients: Iterable[torch.Tensor],
     parameters: Iterable[torch.Tensor],
+    zero_mode,
 ) -> float:
     """Get L2 norm
     Arguments:
@@ -61,7 +62,17 @@ def compute_norm(
     if DTENSOR_SUPPORTED and isinstance(total_norm, DTensor):
         total_norm = total_norm.full_tensor()
 
-    dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=gpc.get_group(ParallelMode.GLOBAL))
+    if gpc.is_using_parallel_mode(zero_mode):
+        dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=gpc.get_group(zero_mode))
+
+    # Need to allreduce(avg) the norms across different ranks because moe params will not be synced during allreduce
+    # model and zero have been reduced!!!
+    if zero_mode == ParallelMode.EXPERT_DATA:
+        pg = gpc.get_group(ParallelMode.EXPERT)
+        scaled_norm = total_norm * 1.0 / float(gpc.get_world_size(ParallelMode.EXPERT))
+        scaled_norm_tensor = torch.tensor(scaled_norm, device=get_current_device(), dtype=torch.float)
+        dist.all_reduce(scaled_norm_tensor, group=pg)
+        total_norm = scaled_norm_tensor.item()
 
     if torch.is_tensor(total_norm):
         total_norm = total_norm.item()
@@ -112,10 +123,14 @@ def __init__(
         # fp16 share mem space with model.FlatParam, fp32 share mem space with optim.param_group
         self._fp16_param_groups = dict()
         self._fp32_param_tensor_groups = dict()
+        self._broadcast_parallel_mode = []
 
         # init fp16 and fp32 params
         for group_idx, param_group in enumerate(self.optim.param_groups):
             group_params = param_group["params"]
+            
+            zero_mode = param_group["optimizer_mode"]
+            self._broadcast_parallel_mode.append(zero_mode)
 
             # fp16 FlatParam storage
             self._fp16_param_groups[group_idx] = group_params
@@ -142,7 +157,7 @@ def _compute_norm_with_fsdp_flatten(self, group_id):
         norm_group = 0
         if len(params) <= 0 or len(gradients) <= 0:
             return norm_group
-        norm_group = compute_norm(gradients=gradients, parameters=params)
+        norm_group = compute_norm(gradients=gradients, parameters=params, zero_mode=self._broadcast_parallel_mode[group_id])
 
         return norm_group