[LLM] First commit the llm deployment code

fastdeploy/model_executor/layers/moe/mm.py

@@ -0,0 +1,273 @@
+"""
+# Copyright (c) 2025 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.
+"""
+
+import os
+import paddle
+from paddle import nn
+from fastdeploy.model_executor.layers.moe.moe import MoELayer
+from fastdeploy.model_executor.layers.utils import get_tensor
+
+
+class TextMoELayer(MoELayer):
+    """
+    MoELayer is a layer that performs MoE (Mixture of Experts) computation.
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        """
+            初始化函数，用于设置类的属性和方法。
+        参数：
+            - args (tuple, optional): 可变长度的位置参数列表，默认为空元组。
+            - kwargs (dict, optional): 关键字参数字典，默认为空字典。
+        返回值：
+            无返回值，直接修改类的属性和方法。
+        """
+        kwargs["moe_tag"] = "Text"
+        super().__init__(*args, **kwargs)
+
+    def load_gate_state_dict(self, state_dict):
+        """
+            加载门状态字典，用于初始化网络参数。
+        将从给定的状态字典中弹出的参数赋值给网络的门参数。
+
+        Args:
+            state_dict (OrderedDict): 包含网络门参数的字典。
+
+        Returns:
+            tuple (list, list): 返回两个列表，分别代表上阶网关投影和下阶投影的参数。
+                每个元素都是一个列表，长度为网络的专家数量。
+        """
+        up_gate_proj_weight = []
+        up_gate_proj_weight_scale = []
+        down_proj_weight = []
+        down_proj_weight_scale = []
+        for j in range(0, self.num_experts):
+            up_gate_proj_weight.append(
+                get_tensor(state_dict.pop(self.ffn1_expert_weight_key.format(j)))
+            )
+            down_proj_weight.append(
+                get_tensor(state_dict.pop(self.ffn2_expert_weight_key.format(j)))
+            )
+        return (
+            up_gate_proj_weight,
+            down_proj_weight,
+            up_gate_proj_weight_scale,
+            down_proj_weight_scale,
+        )
+
+    def load_gate_correction_bias(self, state_dict):
+        """
+            加载网关校正偏置。如果使用了网关校正偏置，则从state_dict中获取相应的张量并设置到网关校正偏置上。
+        参数：
+            state_dict (OrderedDict): 包含模型参数和状态的字典。
+        返回值：
+            无返回值，直接修改了网关校正偏置的值。
+        """
+        if self.moe_config.moe_use_gate_correction_bias:
+            gate_correction_bias_tensor = get_tensor(
+                state_dict[self.gate_correction_bias_key]
+            )
+            self.gate_correction_bias.set_value(
+                gate_correction_bias_tensor[0].unsqueeze(0)
+            )
+
+
+class ImageMoELayer(MoELayer):
+    """
+    MoELayer is a layer that performs MoE (Mixture of Experts) computation.
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        """
+            初始化函数，用于设置类的属性和方法。
+        参数：
+            - args (tuple, optional): 可变长度的位置参数列表，默认为空元组。
+            - kwargs (dict, optional): 关键字参数字典，默认为空字典。
+        返回值：
+            无返回值，直接修改类的属性和方法。
+        """
+        moe_quant_type = os.getenv("ELLM_MM_IMAGE_QUANT_TYPE", None)
+        if moe_quant_type is not None:
+            kwargs["moe_quant_type"] = moe_quant_type
+        kwargs["moe_tag"] = "Image"
+        super().__init__(*args, **kwargs)
+
+    def load_gate_state_dict(self, state_dict):
+        """
+            加载门状态字典。
+        从给定的状态字典中提取并返回两个专家的上下关门投影权重，以及两个专家的下降投影权重。
+        参数：
+            state_dict (OrderedDict): 包含网络参数的有序字典。
+        返回值：
+            tuple (list, list)，分别是两个专家的上下关门投影权重和两个专家的下降投影权重，都是列表类型。
+        """
+        up_gate_proj_weight = []
+        up_gate_proj_weight_scale = []
+        down_proj_weight = []
+        down_proj_weight_scale = []
+        for j in range(self.num_experts, self.num_experts + self.num_experts):
+            up_gate_proj_weight.append(
+                get_tensor(state_dict.pop(self.ffn1_expert_weight_key.format(j)))
+            )
+            down_proj_weight.append(
+                get_tensor(state_dict.pop(self.ffn2_expert_weight_key.format(j)))
+            )
+        return (
+            up_gate_proj_weight,
+            down_proj_weight,
+            up_gate_proj_weight_scale,
+            down_proj_weight_scale,
+        )
+
+    def load_gate_correction_bias(self, state_dict):
+        """
+            加载门级别校正偏置参数，如果使用门级别校正偏置则从state_dict中获取并设置到gate_correction_bias中。
+        参数：
+            state_dict (OrderedDict): 模型的状态字典，包含所有需要被加载的参数。
+        返回值：
+            无返回值，直接修改了gate_correction_bias的值。
+        """
+        if self.moe_config.moe_use_gate_correction_bias:
+            gate_correction_bias_tensor = get_tensor(
+                state_dict[self.gate_correction_bias_key]
+            )
+            self.gate_correction_bias.set_value(
+                gate_correction_bias_tensor[1].unsqueeze(0)
+            )
+
+
+class MultimodalityMoeLayer(nn.Layer):
+    """
+    Multimodality MOE Layer
+    """
+
+    def __init__(
+        self,
+        inference_args,
+        layer_name,
+        layer_idx,
+    ):
+        """
+            初始化一个 MoELayer。
+
+        Args:
+            inference_args (InferenceArgs): 推理参数类，包含了所有必要的配置信息。
+            layer_name (str): 当前 MoE Layer 的名称。
+            layer_idx (int): 当前 MoE Layer 在模型中的索引。
+
+        Returns:
+            None, 无返回值。
+        """
+        super().__init__()
+
+        self.text_moe_layer = TextMoELayer(
+            inference_args=inference_args,
+            moe_config=inference_args.moe_config,
+            layer_name=layer_name + ".text",
+            gate_weight_key=f"ernie.layers.{layer_idx}.mlp.gate.weight",
+            ffn1_expert_weight_key=f"ernie.layers.{layer_idx}.mlp.experts"
+            + ".{}.up_gate_proj.weight",
+            ffn2_expert_weight_key=f"ernie.layers.{layer_idx}.mlp.experts"
+            + ".{}.down_proj.weight",
+            gate_correction_bias_key=f"ernie.layers.{layer_idx}.mlp.moe_statics.e_score_correction_bias",
+            ffn1_bias_key=None,
+            ffn2_bias_key=None,
+            ffn1_shared_weight_key=None,
+            ffn1_shared_bias_key=None,
+            ffn2_shared_weight_key=None,
+            ffn2_shared_bias_key=None,
+            layer_idx=layer_idx,
+        )
+
+        self.image_moe_layer = ImageMoELayer(
+            inference_args=inference_args,
+            moe_config=inference_args.moe_config_1,
+            layer_name=layer_name + ".image",
+            gate_weight_key=f"ernie.layers.{layer_idx}.mlp.gate.weight_1",
+            ffn1_expert_weight_key=f"ernie.layers.{layer_idx}.mlp.experts"
+            + ".{}.up_gate_proj.weight",
+            ffn2_expert_weight_key=f"ernie.layers.{layer_idx}.mlp.experts"
+            + ".{}.down_proj.weight",
+            gate_correction_bias_key=f"ernie.layers.{layer_idx}.mlp.moe_statics.e_score_correction_bias",
+            ffn1_bias_key=None,
+            ffn2_bias_key=None,
+            ffn1_shared_weight_key=None,
+            ffn1_shared_bias_key=None,
+            ffn2_shared_weight_key=None,
+            ffn2_shared_bias_key=None,
+            layer_idx=layer_idx,
+        )
+
+    def load_state_dict(self, state_dict):
+        """
+            加载模型参数。
+        将给定的字典中的参数覆盖到当前模型上，并返回一个新的字典，其中包含未被覆盖的键值对。
+
+        Args:
+            state_dict (dict): 包含了要加载的模型参数的字典。
+
+        Returns:
+            dict: 包含未被覆盖的键值对的字典。
+        """
+        self.text_moe_layer.load_state_dict(state_dict)
+        self.image_moe_layer.load_state_dict(state_dict)
+        state_dict.pop(self.text_moe_layer.gate_correction_bias_key)
+
+    def forward(self, x, **kwargs):
+        """
+            前向计算函数，将输入的张量进行处理并返回结果。
+        该函数接受以下键值对参数：
+            - token_type_ids (Optional, Tensor, default=None): 一个bool型Tensor，用于指定每个元素是否为文本类型（值为0）或图像类型（值为1）。
+                如果未提供此参数，则会引发AssertionError。
+        返回值是一个Tensor，形状与输入相同，表示处理后的结果。
+
+        Args:
+            x (Tensor): 输入张量，形状为[token_num, hidden_size]，其中token_num是序列长度，hidden_size是隐藏状态维度。
+            kwargs (dict, optional): 可选参数字典，默认为None，包含以下键值对：
+                - token_type_ids (Tensor, optional): 一个bool型Tensor，用于指定每个元素是否为文本类型（值为0）或图像类型（值为1），默认为None。
+
+        Returns:
+            Tensor: 一个Tensor，形状与输入相同，表示处理后的结果。
+
+        Raises:
+            AssertionError: 当未提供token_type_ids参数时会引发此错误。
+        """
+        token_type_ids = kwargs.get("token_type_ids", None)
+        assert token_type_ids is not None
+
+        # x.shape is [token_num, hidden_size]
+        fused_moe_out = paddle.zeros_like(x)
+
+        text_mask = token_type_ids == 0  # [token_num]
+        image_mask = token_type_ids == 1
+
+        if text_mask.any():
+            text_out = self.text_moe_layer(x[text_mask])
+            fused_moe_out[text_mask] = text_out
+
+        if image_mask.any():
+            image_out = self.image_moe_layer(x[image_mask])
+            fused_moe_out[image_mask] = image_out
+
+        return fused_moe_out