fix(isp.py): fix isp overlap backward allgather twice when activation ckpt 0.x (#366)

Author: huangting4201

internlm/core/parallel/comm/isp.py

@@ -488,8 +488,8 @@ def _pre_forward_hook_for_module(self, module: nn.Module, *args):  # pylint: dis
         self._wait_handle(module)
 
     def _post_forward_hook_for_module(self, module: nn.Module, *args):  # pylint: disable=W0613
-        self._clear_handle(module)
         if not ((self._module_to_index[module] < self._ckpt_block_num) and self.is_forward is False):
+            self._clear_handle(module)
             self._clear_weight(module)
 
     def _pre_backward_hook_for_module(self, module: nn.Module, *args):  # pylint: disable=W0613

internlm/model/modeling_internlm.py

@@ -138,7 +138,7 @@ def __init__(
             mlp_layer_fusion=mlp_layer_fusion,
             multiple_of=multiple_of,
             # TODO: to support more activation functions
-            activation_type="swiglu" if use_swiglu else "swiglu",
+            activation_type="swiglu" if use_swiglu else "gelu",
         )
 
         self.use_swiglu = use_swiglu

internlm/model/modeling_internlm2.py

@@ -165,7 +165,7 @@ def __init__(
             mlp_layer_fusion=mlp_layer_fusion,
             multiple_of=multiple_of,
             # TODO: to support more activation functions
-            activation_type="swiglu" if use_swiglu else "swiglu",
+            activation_type="swiglu" if use_swiglu else "gelu",
         )
 
         self.use_swiglu = use_swiglu

internlm/model/modeling_llama.py

@@ -157,7 +157,7 @@ def __init__(
             mlp_layer_fusion=mlp_layer_fusion,
             multiple_of=multiple_of,
             # TODO: to support more activation functions
-            activation_type="swiglu" if use_swiglu else "swiglu",
+            activation_type="swiglu" if use_swiglu else "gelu",
         )
 
         self.use_swiglu = use_swiglu

internlm/model/modeling_mixtral.py

@@ -138,7 +138,7 @@ def __init__(
                 mlp_layer_fusion=mlp_layer_fusion,
                 multiple_of=multiple_of,
                 # TODO: to support more activation functions
-                activation_type="swiglu" if use_swiglu else "swiglu",
+                activation_type="swiglu" if use_swiglu else "gelu",
             )
         else:
             # replace mlp by MoE module. The expert in MoE is a FeedForward module.
@@ -156,7 +156,7 @@ def __init__(
                 mlp_layer_fusion=mlp_layer_fusion,
                 multiple_of=multiple_of,
                 # TODO: to support more activation functions
-                activation_type="swiglu" if use_swiglu else "swiglu",
+                activation_type="swiglu" if use_swiglu else "gelu",
             )
 
         self.use_swiglu = use_swiglu

internlm/model/modeling_moe.py

@@ -129,7 +129,7 @@ def __init__(
                 mlp_layer_fusion=mlp_layer_fusion,
                 multiple_of=multiple_of,
                 # TODO: to support more activation functions
-                activation_type="swiglu" if use_swiglu else "swiglu",
+                activation_type="swiglu" if use_swiglu else "gelu",
             )
         else:
             # replace mlp by MoE module. The expert in MoE is a FeedForward module.
@@ -147,7 +147,7 @@ def __init__(
                 mlp_layer_fusion=mlp_layer_fusion,
                 multiple_of=multiple_of,
                 # TODO: to support more activation functions
-                activation_type="swiglu" if use_swiglu else "swiglu",
+                activation_type="swiglu" if use_swiglu else "gelu",
             )
 
         self.use_swiglu = use_swiglu