[Feature] Fused Mixtral support

Author: penPenf28

llm/predict/predictor.py

@@ -118,7 +118,9 @@ class PredictorArgument:
     block_size: int = field(default=64, metadata={"help": "the block size for cache_kvs."})
     cachekv_int8_type: str = field(
         default=None,
-        metadata={"help": "If cachekv_int8_type set as `dynamic`, cache kv would be quantized to int8 dynamically. If cachekv_int8_type set as `static`, cache kv would be quantized to int8 Statically."},
+        metadata={
+            "help": "If cachekv_int8_type set as `dynamic`, cache kv would be quantized to int8 dynamically. If cachekv_int8_type set as `static`, cache kv would be quantized to int8 Statically."
+        },
     )
 
     chat_template: str = field(
@@ -1090,9 +1092,7 @@ def __init__(
         if config.cachekv_int8_type is not None:
             cachekv_dtype = "uint8"
         for i in range(len(self.cache_kvs_shape) // 2):
-            self.cache_kvs["key_caches_{}".format(i)] = paddle.zeros(
-                self.cache_kvs_shape[2 * i], dtype=cachekv_dtype
-            )
+            self.cache_kvs["key_caches_{}".format(i)] = paddle.zeros(self.cache_kvs_shape[2 * i], dtype=cachekv_dtype)
             self.cache_kvs["value_caches_{}".format(i)] = paddle.zeros(
                 self.cache_kvs_shape[2 * i + 1], dtype=cachekv_dtype
             )
@@ -1413,6 +1413,35 @@ def create_predictor(
                     )
                 model.eval()
 
+            elif "mixtral" in config.architectures[0].lower():
+                if predictor_args.block_attn:
+                    config.max_seq_len = predictor_args.total_max_length
+                    config.block_size = predictor_args.block_size
+                    from paddlenlp.experimental.transformers import (
+                        MixtralForCausalLMBlockInferenceModel as MixtralInferenceModel,
+                    )
+
+                    model = MixtralInferenceModel.from_pretrained(
+                        predictor_args.model_name_or_path,
+                        config=config,
+                        dtype=predictor_args.dtype,
+                        tensor_parallel_degree=tensor_parallel_degree,
+                        tensor_parallel_rank=tensor_parallel_rank,
+                    )
+                else:
+                    from paddlenlp.experimental.transformers import (
+                        MixtralForCausalLMInferenceModel as MixtralInferenceModel,
+                    )
+
+                    model = MixtralInferenceModel.from_pretrained(
+                        predictor_args.model_name_or_path,
+                        config=config,
+                        dtype=predictor_args.dtype,
+                        tensor_parallel_degree=tensor_parallel_degree,
+                        tensor_parallel_rank=tensor_parallel_rank,
+                    )
+                model.eval()
+
             elif "opt" in config.architectures[0].lower():
                 if model_args.model_type == "opt-img2txt":
                     # we use opt for img2txt.
@@ -1525,6 +1554,20 @@ def create_predictor(
                 cache_kvs_shape = LlamaInferenceModel.get_cache_kvs_shape(
                     config, predictor_args.batch_size, predictor_args.total_max_length
                 )
+            elif "mixtral" in config.architectures[0].lower():
+                if predictor_args.block_attn:
+                    config.block_size = predictor_args.block_size
+                    config.max_seq_len = predictor_args.total_max_length
+                    from paddlenlp.experimental.transformers import (
+                        MixtralForCausalLMBlockInferenceModel as MixtralInferenceModel,
+                    )
+                else:
+                    from paddlenlp.experimental.transformers import (
+                        MixtralForCausalLMInferenceModel as MixtralInferenceModel,
+                    )
+                cache_kvs_shape = MixtralInferenceModel.get_cache_kvs_shape(
+                    config, predictor_args.batch_size, predictor_args.total_max_length
+                )
             elif "chatglmv2forcausallm" in config.architectures[0].lower():
                 from paddlenlp.experimental.transformers import (
                     ChatGLMv2ForCausalLMInferenceModel,

paddlenlp/experimental/transformers/__init__.py

@@ -18,5 +18,6 @@
 from .fused_transformer_layers import *
 from .gpt import *
 from .llama import *
+from .mixtral import *
 from .opt import *
 from .qwen import *

paddlenlp/experimental/transformers/fused_transformer_layers.py

@@ -18,6 +18,7 @@
 from paddle.framework import LayerHelper, core, in_dynamic_mode
 from paddle.incubate.nn.functional import (
     fused_layer_norm,
+    fused_moe,
     fused_rms_norm,
     masked_multihead_attention,
     variable_length_memory_efficient_attention,
@@ -167,6 +168,7 @@ def __init__(
         linear_bias_attrs=None,
         ffn_ln_scale_attrs=None,
         ffn_ln_bias_attrs=None,
+        gate_weight_attrs=None,
         ffn1_weight_attrs=None,
         ffn1_weight_scale_attrs=None,
         ffn1_bias_attrs=None,
@@ -197,12 +199,15 @@ def __init__(
         kv_num_heads=-1,
         cachekv_int8_type=None,
         rank_id=-1,
+        is_moe=False,
+        moe_every2=False,
+        moe_topk=2,
+        num_experts=1,
     ):
         self.embed_dim = embed_dim
         self.num_heads = num_heads
         if kv_num_heads > 0:
             self.kv_num_heads = kv_num_heads
-            assert nranks == 1, "nranks should be 1 for kv_num_heads > 0"
         else:
             self.kv_num_heads = num_heads
         self.dim_feedforward = dim_feedforward
@@ -222,6 +227,7 @@ def __init__(
         self.linear_bias_attrs = linear_bias_attrs
         self.ffn_ln_scale_attrs = ffn_ln_scale_attrs
         self.ffn_ln_bias_attrs = ffn_ln_bias_attrs
+        self.gate_weight_attrs = gate_weight_attrs
         self.ffn1_weight_attrs = ffn1_weight_attrs
         self.ffn1_weight_scale_attrs = ffn1_weight_scale_attrs
         self.ffn1_bias_attrs = ffn1_bias_attrs
@@ -255,6 +261,10 @@ def __init__(
         self.rank_id = rank_id
         self.trans_qkvw = trans_qkvw
         self.ring_id = ring_id
+        self.is_moe = is_moe
+        self.moe_every2 = moe_every2
+        self.moe_topk = moe_topk
+        self.num_experts = num_experts
 
 
 class FusedMultiTransformerBase(Layer):
@@ -294,6 +304,10 @@ 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"
 
+        self._is_moe = config.is_moe
+        self._moe_every2 = config.moe_every2
+        self._moe_topk = config.moe_topk
+
         # tensor model parallel
         if config.nranks > 1:
             assert config.ring_id != -1
@@ -316,6 +330,7 @@ def __init__(self, config: FusedMultiTransformerConfig):
         self.qkv_weights, self.qkv_biases = [], []
         self.linear_weights, self.linear_biases = [], []
         self.ffn_ln_scales, self.ffn_ln_biases = [], []
+        self.gate_weights = []
         self.ffn1_weights, self.ffn1_biases = [], []
         self.ffn2_weights, self.ffn2_biases = [], []
         self.cache_k_scales, self.cache_v_scales = [], []
@@ -327,6 +342,7 @@ def __init__(self, config: FusedMultiTransformerConfig):
             qkv_weight_attr = self.get_attr(config.qkv_weight_attrs, i)
 
             qkv_bias_attr = self.get_attr(config.qkv_bias_attrs, i)
+            gate_weight_attr = self.get_attr(config.gate_weight_attrs, i)
             linear_weight_attr = self.get_attr(config.linear_weight_attrs, i)
             linear_bias_attr = self.get_attr(config.linear_bias_attrs, i)
 
@@ -407,37 +423,99 @@ def __init__(self, config: FusedMultiTransformerConfig):
                     dtype=self._norm_weight_dtype,
                 )
 
-            ffn1_weight = self.create_parameter(
-                shape=self.ffn1_weight_shape,
-                attr=ffn1_weight_attr,
-                dtype=self.create_params_type,
-                is_bias=False,
-            )
+            gate_weight = None
+            if config.is_moe is True and ((config.moe_every2 is True and i % 2 == 1) or config.moe_every2 is False):
+                gate_weight = self.create_parameter(
+                    shape=[config.embed_dim, config.num_experts],
+                    attr=gate_weight_attr,
+                    dtype="float32",
+                    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],
+                    attr=ffn1_weight_attr,
+                    dtype=self.create_params_type,
+                    is_bias=False,
+                )
+            else:
+                ffn1_weight = self.create_parameter(
+                    shape=self.ffn1_weight_shape,
+                    attr=ffn1_weight_attr,
+                    dtype=self.create_params_type,
+                    is_bias=False,
+                )
 
             ffn1_bias = None
             if ffn1_bias_attr:
-                ffn1_bias = self.create_parameter(
-                    shape=[dim_feedforward * 2] if config.activation.endswith("glu") else [dim_feedforward],
-                    attr=ffn1_bias_attr,
-                    dtype=self._dtype,
-                    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
+                ):
+                    ffn1_bias = self.create_parameter(
+                        shape=[config.num_experts, self.dim_feedforward * 2]
+                        if self.activation.endswith("glu")
+                        else [config.num_experts, self.dim_feedforward],
+                        attr=ffn1_bias_attr,
+                        dtype=self._dtype,
+                        is_bias=True,
+                    )
+                else:
+                    ffn1_bias = self.create_parameter(
+                        shape=[dim_feedforward * 2] if self.activation.endswith("glu") else [dim_feedforward],
+                        attr=ffn1_bias_attr,
+                        dtype=self._dtype,
+                        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):
+                ffn2_weight = self.create_parameter(
+                    shape=[config.num_experts, self.dim_feedforward, self.embed_dim],
+                    attr=ffn2_weight_attr,
+                    dtype=self.create_params_type,
+                    is_bias=False,
+                )
+            else:
+                ffn2_weight = self.create_parameter(
+                    shape=self.ffn2_weight_shape,
+                    attr=ffn2_weight_attr,
+                    dtype=self.create_params_type,
+                    is_bias=False,
                 )
-
-            ffn2_weight = self.create_parameter(
-                shape=self.ffn2_weight_shape,
-                attr=ffn2_weight_attr,
-                dtype=self.create_params_type,
-                is_bias=False,
-            )
 
             ffn2_bias = None
             if ffn2_bias_attr:
-                ffn2_bias = self.create_parameter(
-                    shape=[config.embed_dim],
-                    attr=ffn2_bias_attr,
-                    dtype=self._dtype,
-                    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
+                ):
+                    ffn2_bias = self.create_parameter(
+                        shape=[config.num_experts, config.embed_dim],
+                        attr=ffn2_bias_attr,
+                        dtype=self._dtype,
+                        is_bias=True,
+                    )
+                else:
+                    ffn2_bias = self.create_parameter(
+                        shape=[config.embed_dim],
+                        attr=ffn2_bias_attr,
+                        dtype=self._dtype,
+                        is_bias=True,
+                    )
 
             cache_k_scale = None
             if cache_k_scale_attr:
@@ -495,6 +573,8 @@ def __init__(self, config: FusedMultiTransformerConfig):
 
             self.ffn_ln_scales.append(ffn_ln_scale)
             self.ffn_ln_biases.append(ffn_ln_bias)
+            if gate_weight is not None:
+                self.gate_weights.append(gate_weight)
             self.ffn1_weights.append(ffn1_weight)
             self.ffn1_biases.append(ffn1_bias)
             self.ffn2_weights.append(ffn2_weight)
@@ -713,6 +793,28 @@ def compute_ffn_layernorm(self, out_linear_out, residual_input, i):
 
         return tmp_out, residual_input
 
+    def compute_fused_moe(self, tmp_out, i):
+        # todo[xinhw]: make bias optional
+        if self.ffn1_biases[i] is None:
+            shape1 = paddle.to_tensor([self.ffn1_weights[i].shape[0], 1, self.dim_feedforward * 2])
+            self.ffn1_biases[i] = paddle.zeros(shape1)
+        if self.ffn2_biases[i] is None:
+            shape2 = paddle.to_tensor([self.ffn1_weights[i].shape[0], 1, self.embed_dim])
+            self.ffn2_biases[i] = paddle.zeros(shape2)
+        fused_moe_out = fused_moe(
+            tmp_out,
+            self.gate_weights[i],
+            self.ffn1_weights[i],
+            self.ffn1_biases[i],
+            self.ffn2_weights[i],
+            self.ffn2_biases[i],
+            None,
+            None,
+            "None",
+            self._moe_topk,
+        )
+        return fused_moe_out
+
     def compute_activation(self, ffn1_out, i):
         return fused_act_bias_wrapper(ffn1_out, self.ffn1_biases[i], act_method=self.activation)
 
@@ -854,12 +956,17 @@ def forward(
             # ffn layernorm
             tmp_out, residual_input = self.compute_ffn_layernorm(out_linear_out, residual_input, i)
 
-            # ffn1 matmul
-            ffn1_out = self.compute_ffn1(tmp_out, i)
-            ffn1_out = self.compute_activation(ffn1_out, i)
+            if self._is_moe is True and ((self._moe_every2 is True and i % 2 == 1) or self._moe_every2 is False):
+                # fused moe
+                ffn2_out = self.compute_fused_moe(tmp_out, i)
+
+            else:
+                # ffn1 matmul
+                ffn1_out = self.compute_ffn1(tmp_out, i)
+                ffn1_out = self.compute_activation(ffn1_out, i)
 
-            # ffn2 matmul
-            ffn2_out = self.compute_ffn2(ffn1_out, i)
+                # ffn2 matmul
+                ffn2_out = self.compute_ffn2(ffn1_out, i)
 
             # all_reduce
             if self.nranks > 1:

paddlenlp/experimental/transformers/mixtral/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .modeling import *