From 1aed8ca42f64ee8c74f5c8589fd58fc1967b5db2 Mon Sep 17 00:00:00 2001
From: yingtongxiong <974106207@qq.com>
Date: Tue, 3 Dec 2024 16:39:31 +0800
Subject: [PATCH 1/3] refactor mlp and add warning

---
 internlm/core/context/parallel_context.py |  6 ++++++
 internlm/initialize/launch.py             |  8 +++++---
 internlm/model/modules/mlp.py             | 24 ++++++++++++-----------
 3 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/internlm/core/context/parallel_context.py b/internlm/core/context/parallel_context.py
index 989b1c00f..e0b6f2a9b 100644
--- a/internlm/core/context/parallel_context.py
+++ b/internlm/core/context/parallel_context.py
@@ -78,6 +78,12 @@ def update(self, config):
             self._add_item(k, v)
         return self
 
+    def __delattr__(self, key):
+        if key in self:
+            super().__delitem__(key)
+        else:
+            raise AttributeError(f"{key} does not exist")
+
     @staticmethod
     def from_file(filename: str):
         """Reads a python file and constructs a corresponding :class:`Config` object.
diff --git a/internlm/initialize/launch.py b/internlm/initialize/launch.py
index 1ac8ef31d..cfd557f93 100644
--- a/internlm/initialize/launch.py
+++ b/internlm/initialize/launch.py
@@ -358,6 +358,9 @@ def args_sanity_check():
                 "Please make sure you are using flash attention in cuda device."
             )
 
+    if "mlp_layer_fusion" not in model:
+        model._add_item("mlp_layer_fusion", False)
+
     if "MoE" in gpc.config.get("model_type", ModelType.INTERNLM.name):
         if "num_experts" not in model:
             model._add_item("num_experts", 1)
@@ -371,9 +374,8 @@ def args_sanity_check():
             model._add_item("moe_type", "GShard")
         if "moe_layer_kwargs" not in model:
             model.moe_layer_kwargs = {}
-
-    if "mlp_layer_fusion" not in model:
-        model._add_item("mlp_layer_fusion", False)
+        if model.mlp_layer_fusion is False:
+            logger.warning("The config 'mlp_layer_fusion' is False, we recommend it should be set True when use MoE.")
 
     # qk_interleaved config
     if "qk_interleaved" not in gpc.config.model:
diff --git a/internlm/model/modules/mlp.py b/internlm/model/modules/mlp.py
index 6e74d6b6f..8c1dee975 100644
--- a/internlm/model/modules/mlp.py
+++ b/internlm/model/modules/mlp.py
@@ -105,21 +105,23 @@ def __init__(
             self.w3 = new_linear(
                 "w3", in_features, hidden_features, bias, device=device, dtype=dtype, is_expert=is_expert
             )
+        
+        if self.activation_type is ActivationType.swiglu.name:
+            self.activation_fn = Silu
+        else:
+            self.activation_fn = Gelu
 
     def forward(self, x):
-        if not self.mlp_layer_fusion:
-            w1_o = self.w1(x)
-            w3_o = self.w3(x)
-        else:
-            fussed_out = self.fused_w1_w3(x)
-            w1_o, w3_o = torch.split(fussed_out, fussed_out.shape[-1] // 2, dim=-1)
 
-        if self.activation_type is ActivationType.swiglu.name:
-            out = self.w2(Silu(w1_o, w3_o))
-        else:
-            out = self.w2(Gelu(w1_o, w3_o))
+        if self.mlp_layer_fusion:
+            fused_out = self.fused_w1_w3(x)
+            w1_o, w3_o = torch.split(fused_out, fused_out.shape[-1] // 2, dim=-1)
+            return self.w2(self.activation_fn(w1_o, w3_o))
+
+        w1_o = self.w1(x)
+        w3_o = self.w3(x)
+        return self.w2(self.activation_fn(w1_o, w3_o))
 
-        return out
 
 
 class GroupedFeedForward(nn.Module):

From 2274aef2ce121505496f703505947a963605a36f Mon Sep 17 00:00:00 2001
From: yingtongxiong <974106207@qq.com>
Date: Tue, 3 Dec 2024 16:41:20 +0800
Subject: [PATCH 2/3] fix lint

---
 internlm/model/modules/mlp.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/internlm/model/modules/mlp.py b/internlm/model/modules/mlp.py
index 8c1dee975..8826b8cbb 100644
--- a/internlm/model/modules/mlp.py
+++ b/internlm/model/modules/mlp.py
@@ -105,7 +105,7 @@ def __init__(
             self.w3 = new_linear(
                 "w3", in_features, hidden_features, bias, device=device, dtype=dtype, is_expert=is_expert
             )
-        
+
         if self.activation_type is ActivationType.swiglu.name:
             self.activation_fn = Silu
         else:
@@ -123,7 +123,6 @@ def forward(self, x):
         return self.w2(self.activation_fn(w1_o, w3_o))
 
 
-
 class GroupedFeedForward(nn.Module):
     """
     Base FeedForward in flash implementation.

From ffae58815899944d72b58b9332485b70baecf890 Mon Sep 17 00:00:00 2001
From: yingtongxiong <974106207@qq.com>
Date: Tue, 3 Dec 2024 17:22:37 +0800
Subject: [PATCH 3/3] refatcor mlp.py

---
 internlm/model/modules/mlp.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/internlm/model/modules/mlp.py b/internlm/model/modules/mlp.py
index 8826b8cbb..4850eb590 100644
--- a/internlm/model/modules/mlp.py
+++ b/internlm/model/modules/mlp.py
@@ -112,14 +112,12 @@ def __init__(
             self.activation_fn = Gelu
 
     def forward(self, x):
-
         if self.mlp_layer_fusion:
             fused_out = self.fused_w1_w3(x)
             w1_o, w3_o = torch.split(fused_out, fused_out.shape[-1] // 2, dim=-1)
-            return self.w2(self.activation_fn(w1_o, w3_o))
-
-        w1_o = self.w1(x)
-        w3_o = self.w3(x)
+        else:
+            w1_o = self.w1(x)
+            w3_o = self.w3(x)
         return self.w2(self.activation_fn(w1_o, w3_o))