kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity
a75618df24
TensorRT-LLMs/tensorrt_llm/layers/moe.py
Kaiyu Xie a75618df24
Update TensorRT-LLM (#667) 
* Update TensorRT-LLM

---------

Co-authored-by: 0xymoro <jerrymeng100@gmail.com>
Co-authored-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
2023-12-15 22:14:51 +08:00

332 lines
14 KiB
Python
Raw Blame History

# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 dataclasses import dataclass
from enum import IntEnum
from typing import List
import numpy as np
import tensorrt as trt
from tensorrt_llm._utils import str_dtype_to_trt
from .._common import default_trtnet
from ..functional import _create_tensor, allreduce, cast, split
from ..module import Module
from ..parameter import Parameter
from ..plugin import TRT_LLM_PLUGIN_NAMESPACE
from ..quantization import QuantMode
from .linear import RowLinear
activation_str_to_int_map = {
# [WARNING] Keep the below in sync with cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels.h
"gelu": 0,
"gelu_new": 0,
"relu": 1,
"silu": 2,
"swiglu": 3,
"geglu": 4,
"identity": 5,
}
@dataclass
class MoeConfig:
# [WARNING] Keep the below in sync with cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.h
class ParallelismMode(IntEnum):
NONE = 0
EXPERT_PARALLEL = 1
TENSOR_PARALLEL = 2
class ExpertScaleNormalizationMode(IntEnum):
NONE = 0
RENORMALIZE = 1
num_experts: int = 0
top_k: int = 0
tp_mode: ParallelismMode = ParallelismMode.TENSOR_PARALLEL
normalization_mode: ExpertScaleNormalizationMode = ExpertScaleNormalizationMode.RENORMALIZE
def validate(self) -> "MoeConfig":
if (self.num_experts == 0) != (self.top_k == 0):
raise ValueError(
"Both or neither MoeConfig's num_experts and top_k must be set to 0"
)
return self
def has_moe(self) -> bool:
return self.num_experts > 1
def is_gated_activation(activation_str):
return activation_str in ("swiglu", "geglu")
def _moe_plugin(moe_config,
hidden_states,
routing,
finished,
expert_weight_1,
expert_weight_2,
expert_bias_1,
expert_bias_2,
expert_scale_1,
expert_scale_2,
hidden_size,
ffn_hidden_size,
act_fn,
dtype,
weight_dtype,
quant_mode=QuantMode(0),
tp_size=1,
tp_rank=0):
if isinstance(dtype, str):
dtype = str_dtype_to_trt(dtype)
# Create the plugin with our required state
num_experts = moe_config.num_experts
# We pass the full number of experts (not divided by tp_size) even for EP mode
p_num_experts = trt.PluginField("number_of_experts",
np.array(num_experts, dtype=np.int32),
trt.PluginFieldType.INT32)
p_top_k = trt.PluginField("top_k", np.array(moe_config.top_k,
dtype=np.int32),
trt.PluginFieldType.INT32)
p_expert_hidden_size = trt.PluginField(
"expert_hidden_size", np.array(hidden_size, dtype=np.int32),
trt.PluginFieldType.INT32)
p_expert_inter_size = trt.PluginField(
"expert_inter_size", np.array(ffn_hidden_size, dtype=np.int32),
trt.PluginFieldType.INT32)
p_activation_type = trt.PluginField(
"activation_type",
np.array(activation_str_to_int_map[act_fn], dtype=np.int32),
trt.PluginFieldType.INT32)
p_type_id = trt.PluginField("type_id", np.array([int(dtype)],
dtype=np.int32),
trt.PluginFieldType.INT32)
p_weight_type_id = trt.PluginField(
"weight_type_id", np.array([int(weight_dtype)], dtype=np.int32),
trt.PluginFieldType.INT32)
p_quant_mode = trt.PluginField("quant_mode",
np.array([int(quant_mode)], dtype=np.int32),
trt.PluginFieldType.INT32)
p_use_finished = trt.PluginField(
"use_finished", np.array([int(finished is not None)], dtype=np.int32),
trt.PluginFieldType.INT32)
p_use_bias = trt.PluginField(
"use_bias", np.array([int(expert_bias_1 is not None)], dtype=np.int32),
trt.PluginFieldType.INT32)
p_tp_size = trt.PluginField("tp_size", np.array(tp_size, dtype=np.int32),
trt.PluginFieldType.INT32)
p_tp_rank = trt.PluginField("tp_rank", np.array(tp_rank, dtype=np.int32),
trt.PluginFieldType.INT32)
p_parallelism_mode = trt.PluginField(
"parallelism_mode", np.array(moe_config.tp_mode, dtype=np.int32),
trt.PluginFieldType.INT32)
p_normalization_mode = trt.PluginField(
"normalization_mode",
np.array(moe_config.normalization_mode, dtype=np.int32),
trt.PluginFieldType.INT32)
pfc = trt.PluginFieldCollection([
p_num_experts, p_top_k, p_expert_hidden_size, p_expert_inter_size,
p_activation_type, p_type_id, p_weight_type_id, p_quant_mode,
p_use_finished, p_use_bias, p_tp_size, p_tp_rank, p_parallelism_mode,
p_normalization_mode
])
# Create the plugin with our constant inputs to the constructor
plugin_creator = trt.get_plugin_registry().get_plugin_creator(
'MixtureOfExperts', '1', TRT_LLM_PLUGIN_NAMESPACE)
assert plugin_creator is not None
moe_plugin = plugin_creator.create_plugin("mixture_of_experts", pfc)
# Instantiate the plugin with our specific inputs
plugin_inputs = [
hidden_states, routing, expert_weight_1.value, expert_weight_2.value
]
if expert_bias_1:
assert expert_bias_2
plugin_inputs += [expert_bias_1.value, expert_bias_2.value]
if finished is not None:
plugin_inputs += [finished]
# Add conditional inputs
if expert_scale_1 is not None:
assert expert_scale_2
plugin_inputs += [expert_scale_1.value, expert_scale_2.value]
plugin_inputs = [i.trt_tensor for i in plugin_inputs]
layer = default_trtnet().add_plugin_v2(plugin_inputs, moe_plugin)
for ii in range(layer.num_inputs):
if layer.get_input(ii).dtype == str_dtype_to_trt("int8"):
layer.get_input(ii).set_dynamic_range(-127, 127)
output = _create_tensor(layer.get_output(0), layer)
return output
class MixtureOfExperts(Module):
def __init__(self,
moe_config: MoeConfig,
hidden_size: int,
ffn_hidden_size: int,
hidden_act: str,
bias: bool = True,
dtype=None,
tp_group: List[int] = None,
tp_size: int = 1,
tp_rank: int = 0,
instance_id: int = 0,
quant_mode=QuantMode(0)):
super().__init__()
self.moe_config = moe_config
self.num_experts = moe_config.num_experts
self.top_k = moe_config.top_k
self.hidden_act = hidden_act
self.hidden_size = hidden_size
self.ffn_hidden_size = ffn_hidden_size
self.dtype = dtype
self.weight_dtype = dtype
self.tp_group = tp_group
self.tp_size = tp_size
self.tp_rank = tp_rank
self.instance_id = instance_id
self.quant_mode = quant_mode
self.experts_per_node = self.num_experts
if moe_config.tp_mode == MoeConfig.ParallelismMode.EXPERT_PARALLEL:
if self.num_experts % self.tp_size != 0:
raise ValueError(
f"MixtureOfExperts - Number of experts {self.num_experts} is not a multiple of EP size {self.tp_size}"
)
self.experts_per_node = self.experts_per_node // tp_size
elif moe_config.tp_mode == MoeConfig.ParallelismMode.TENSOR_PARALLEL:
if self.ffn_hidden_size % self.tp_size != 0:
raise ValueError(
f"MixtureOfExperts - FFN Hidden Size {self.ffn_hidden_size} is not a multiple of TP size {self.tp_size}"
)
self.ffn_hidden_size = self.ffn_hidden_size // tp_size
if quant_mode.is_weight_only():
self.weight_dtype = trt.int8
# TODO: benchmark the router and check best TP configuration
# Since output dimension is usually low (in the order of 10s), we split on input dim for the moment
# Maybe no TP at all is even more efficient
self.router = RowLinear(
hidden_size,
self.num_experts,
bias=False,
dtype=trt.
float32, # Routing is sensitive since it conditions what experts are used
tp_group=tp_group,
tp_size=tp_size,
strict_dtype=True,
)
# Note we use horizontal fusion for gated activation to do the operation in one GEMM invocation
# The left matrix is a linear projection (no activation applied)
# The right matrix is the gating value (activation applied)
# The naming convention is the inverse of GatedMLP, but the same as `tensorrt_llm/functional.py`
expert_1_out_size = self.ffn_hidden_size * 2 if is_gated_activation(
hidden_act) else self.ffn_hidden_size
expert_1_shape = (self.experts_per_node, expert_1_out_size, hidden_size)
expert_2_shape = (self.experts_per_node, hidden_size,
self.ffn_hidden_size)
if quant_mode.is_weight_only():
bytes_per_col_scale = 2 if quant_mode.is_int4_weight_only() else 1
# We use a different shape here because the quantized weights have their own layout
expert_1_shape = (self.experts_per_node, hidden_size,
expert_1_out_size // bytes_per_col_scale)
expert_2_shape = (self.experts_per_node, self.ffn_hidden_size,
hidden_size // bytes_per_col_scale)
self.experts_scale_1 = Parameter(shape=(self.experts_per_node,
expert_1_out_size),
dtype=dtype)
self.experts_scale_2 = Parameter(shape=(self.experts_per_node,
hidden_size),
dtype=dtype)
else:
self.register_parameter('experts_scale_1', None)
self.register_parameter('experts_scale_2', None)
self.experts_weight_1 = Parameter(shape=expert_1_shape,
dtype=self.weight_dtype)
self.experts_weight_2 = Parameter(shape=expert_2_shape,
dtype=self.weight_dtype)
# Note: the bias uses dtype NOT weight_dtype, i.e. it is not quantized
if bias:
self.experts_bias_1 = Parameter(shape=(self.experts_per_node,
expert_1_out_size),
dtype=dtype)
self.experts_bias_2 = Parameter(shape=(self.experts_per_node,
hidden_size),
dtype=dtype)
else:
self.register_parameter('experts_bias_1', None)
self.register_parameter('experts_bias_2', None)
def forward(self,
hidden_states,
finished=None,
workspace=None,
lora_layer_params=None):
assert lora_layer_params is None, "LoRA + MoE is not supported for the moment"
routing_input = cast(hidden_states, trt.float32)
if self.tp_size > 1:
routing_input = split(routing_input,
self.router.in_features,
dim=-1)[self.tp_rank]
routing = self.router(routing_input)
output = _moe_plugin(self.moe_config,
hidden_states,
routing,
expert_weight_1=self.experts_weight_1,
expert_weight_2=self.experts_weight_2,
expert_bias_1=self.experts_bias_1,
expert_bias_2=self.experts_bias_2,
expert_scale_1=self.experts_scale_1,
expert_scale_2=self.experts_scale_2,
finished=finished,
hidden_size=self.hidden_size,
ffn_hidden_size=self.ffn_hidden_size,
act_fn=self.hidden_act,
dtype=self.dtype,
weight_dtype=self.weight_dtype,
quant_mode=self.quant_mode,
tp_size=self.tp_size,
tp_rank=self.tp_rank)
if self.tp_size > 1 and self.tp_group is not None and self.moe_config.tp_mode != MoeConfig.ParallelismMode.NONE:
output = allreduce(output,
self.tp_group,
workspace=workspace,
instance_id=self.instance_id)
return output
MOE = MixtureOfExperts
Reference in New Issue View Git Blame Copy Permalink