fix(isp.py): fix isp overlap backward allgather twice when activation…

… ckpt 0.x (#366)
InternLM · Nov 29, 2024 · 4a6b453 · 4a6b453
1 parent 86c3b6c
commit 4a6b453
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 8 deletions.
diff --git a/internlm/core/parallel/comm/isp.py b/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

diff --git a/internlm/model/modeling_internlm.py b/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

diff --git a/internlm/model/modeling_internlm2.py b/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

diff --git a/internlm/model/modeling_llama.py b/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

diff --git a/internlm/model/modeling_mixtral.py b/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

diff --git a/internlm/model/modeling_moe.py b/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