[Refactor] refine code

Author: penPenf28

paddlenlp/experimental/transformers/fused_transformer_layers.py

@@ -316,15 +316,6 @@ def __init__(self, config: FusedMultiTransformerConfig):
         self.head_dim = config.embed_dim // config.num_heads
         assert self.head_dim * config.num_heads == config.embed_dim, "embed_dim must be divisible by num_heads"
 
-        if config.is_moe:
-            self._is_moe = config.is_moe
-            self._moe_every2 = config.moe_every2
-            self._moe_topk = config.moe_topk
-        else:
-            self._is_moe = False
-            self._moe_every2 = False
-            self._moe_topk = 1
-
         # tensor model parallel
         if config.nranks > 1:
             assert config.ring_id != -1
@@ -449,24 +440,10 @@ def __init__(self, config: FusedMultiTransformerConfig):
                     is_bias=False,
                     default_initializer=paddle.nn.initializer.Constant(0),
                 )
-            else:
-                gate_weight = self.create_parameter(
-                    shape=[1],
-                    attr=gate_weight_attr,
-                    dtype="float32",
-                    is_bias=False,
-                    default_initializer=paddle.nn.initializer.Constant(0),
-                )
-
-            if config.is_moe is False:
-                gate_weight = None
-                self.gate_weights = None
 
             if config.is_moe is True and ((config.moe_every2 is True and i % 2 == 1) or config.moe_every2 is False):
                 ffn1_weight = self.create_parameter(
-                    shape=[config.num_experts, self.embed_dim, self.dim_feedforward * 2]
-                    if self.activation.endswith("glu")
-                    else [config.num_experts, self.embed_dim, self.dim_feedforward],
+                    shape=self.moe_ffn1_weight_shape,
                     attr=ffn1_weight_attr,
                     dtype=self.create_params_type,
                     is_bias=False,
@@ -481,8 +458,8 @@ def __init__(self, config: FusedMultiTransformerConfig):
 
             ffn1_bias = None
             if ffn1_bias_attr:
-                if config.is_moe is True and (
-                    (config.moe_every2 is True and i % 2 == 1) or config.moe_every2 is False
+                if self.config.is_moe is True and (
+                    (self.config.moe_every2 is True and i % 2 == 1) or self.config.moe_every2 is False
                 ):
                     ffn1_bias = self.create_parameter(
                         shape=[config.num_experts, self.dim_feedforward * 2]
@@ -500,9 +477,11 @@ def __init__(self, config: FusedMultiTransformerConfig):
                         is_bias=True,
                     )
 
-            if config.is_moe is True and ((config.moe_every2 is True and i % 2 == 1) or config.moe_every2 is False):
+            if self.config.is_moe is True and (
+                (self.config.moe_every2 is True and i % 2 == 1) or self.config.moe_every2 is False
+            ):
                 ffn2_weight = self.create_parameter(
-                    shape=[config.num_experts, self.dim_feedforward, self.embed_dim],
+                    shape=self.moe_ffn2_weight_shape,
                     attr=ffn2_weight_attr,
                     dtype=self.create_params_type,
                     is_bias=False,
@@ -657,6 +636,14 @@ def init_weight_shape(self, config):
         )
         self.ffn2_weight_shape = [self.dim_feedforward, self.embed_dim]
 
+        if self.config.is_moe is True:
+            self.moe_ffn1_weight_shape = (
+                [config.num_experts, self.embed_dim, self.dim_feedforward * 2]
+                if self.activation.endswith("glu")
+                else [config.num_experts, self.embed_dim, self.dim_feedforward]
+            )
+            self.moe_ffn2_weight_shape = [config.num_experts, self.dim_feedforward, self.embed_dim]
+
     def get_weight_create_dype(self):
         return self._dtype
 
@@ -830,7 +817,7 @@ def compute_fused_moe(self, tmp_out, i):
             None,
             None,
             "None",
-            self._moe_topk,
+            self.config.moe_topk,
         )
         return fused_moe_out
 
@@ -975,7 +962,9 @@ def forward(
             # ffn layernorm
             tmp_out, residual_input = self.compute_ffn_layernorm(out_linear_out, residual_input, i)
 
-            if self._is_moe is True and ((self._moe_every2 is True and i % 2 == 1) or self._moe_every2 is False):
+            if self.config.is_moe is True and (
+                (self.config.moe_every2 is True and i % 2 == 1) or self.config.moe_every2 is False
+            ):
                 # fused moe
                 ffn2_out = self.compute_fused_moe(tmp_out, i)
 

paddlenlp/experimental/transformers/mixtral/modeling.py

@@ -774,7 +774,6 @@ def set_state_dict(self, state_dict):
                 act_scale_map_dict = scale_map_dict["act_scale"]
                 weight_scale_map_dict = scale_map_dict["weight_scale"]
                 cache_scale_map_dict = scale_map_dict["cachekv_scale"]
-                # TODO(RichardWooSJTU): support multi-cards
 
                 act_scale_json_path = os.path.join(self.quant_model_path, "act_scales.json")
                 weight_scale_json_path = os.path.join(self.quant_model_path, "weight_scales.json")