[Feature] fused mixtral wint4 (PaddlePaddle#9013)

* [Feature] fused mixtral wint4

* [Refactor] refine code

Author: penPenf28

paddlenlp/experimental/transformers/fused_transformer_layers.py

@@ -1137,6 +1137,10 @@ def init_weight_shape(self, config):
             )
             self.moe_ffn2_weight_shape = [self.config.moe_config.num_experts, self.dim_feedforward, self.embed_dim]
 
+            if config.quant_type == "weight_only_int4":
+                self.moe_ffn1_weight_shape[2] //= 2
+                self.moe_ffn2_weight_shape[2] //= 2
+
     def compute_qkv_linear(self, ln_out, i):
         return weight_only_linear(
             ln_out,

paddlenlp/experimental/transformers/mixtral/modeling.py

@@ -642,7 +642,11 @@ def set_state_dict(self, state_dict):
                         ffn1_weight_tensor[i], algo=self.quant_algo
                     )
                     ffn1_quanted_weight_list.append(
-                        ffn1_quanted_weight_list_i.reshape([self.transformer_block.config.embed_dim, -1])
+                        ffn1_quanted_weight_list_i.reshape(
+                            [self.transformer_block.embed_dim, self.transformer_block.dim_feedforward * 2]
+                            if self.quant_type == "weight_only_int8"
+                            else [self.transformer_block.embed_dim, self.transformer_block.dim_feedforward]
+                        )
                     )
                     ffn1_quanted_weight_scale.append(ffn1_quanted_weight_scale_i)
                 ffn1_quanted_weight_tensor = paddle.to_tensor(ffn1_quanted_weight_list)
@@ -677,7 +681,11 @@ def set_state_dict(self, state_dict):
                         ffn2_weight_tensor[i], algo=self.quant_algo
                     )
                     ffn2_quanted_weight_list.append(
-                        ffn2_quanted_weight_list_i.reshape([-1, self.transformer_block.config.embed_dim])
+                        ffn2_quanted_weight_list_i.reshape(
+                            [self.transformer_block.dim_feedforward, self.transformer_block.embed_dim]
+                            if self.quant_type == "weight_only_int8"
+                            else [self.transformer_block.dim_feedforward, self.transformer_block.embed_dim // 2]
+                        )
                     )
                     ffn2_quanted_weight_scale.append(ffn2_quanted_weight_scale_i)
                 ffn2_quanted_weight_tensor = paddle.to_tensor(ffn2_quanted_weight_list)