mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-10-05 16:48:03 +08:00
274 lines
10 KiB
Python
274 lines
10 KiB
Python
"""
|
||
# 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
|