chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Compat with latest VLLM 0.4.2 main + fork.number rename + Flashinfer 0.0.4 (#380)

Browse Source 
Co-authored-by: ZX <zx@lbx.dev>
Co-authored-by: ZhouXingg <165115237+ZhouXingg@users.noreply.github.com>
This commit is contained in:
Qubitium
2024-05-11 16:37:49 -07:00
committed by GitHub
co-authored by ZX ZhouXingg
parent a511a2d089
commit 33b242df30
20 changed files with 611 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/mixtral.py
+23 -22
View File
@@ -1,7 +1,7 @@
# Adapted from
# https://github.com/vllm-project/vllm/blob/d0215a58e78572d91dadafe9d832a2db89b09a13/vllm/model_executor/models/mixtral.py#L1
"""Inference-only Mixtral model."""
from typing import List, Optional, Tuple
from typing import Optional
import numpy as np
import torch
@@ -10,24 +10,25 @@ from torch import nn
from transformers import MixtralConfig
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
LinearMethodBase,
QKVParallelLinear,
ReplicatedLinear,
RowParallelLinear,
)
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig)
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from vllm.model_executor.parallel_utils.communication_op import (
from vllm.distributed import (
tensor_model_parallel_all_reduce,
)
from vllm.model_executor.parallel_utils.parallel_state import (
from vllm.distributed import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
)
from vllm.model_executor.weight_utils import (
from sglang.srt.weight_utils import (
default_weight_loader,
hf_model_weights_iterator,
)
@@ -43,7 +44,7 @@ class MixtralMLP(nn.Module):
num_experts: int,
hidden_size: int,
intermediate_size: int,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.num_experts = num_experts
@@ -51,13 +52,13 @@ class MixtralMLP(nn.Module):
self.hidden_dim = hidden_size
self.w1 = ReplicatedLinear(
self.hidden_dim, self.ffn_dim, bias=False, linear_method=linear_method
self.hidden_dim, self.ffn_dim, bias=False, quant_config=quant_config
)
self.w2 = ReplicatedLinear(
self.ffn_dim, self.hidden_dim, bias=False, linear_method=linear_method
self.ffn_dim, self.hidden_dim, bias=False, quant_config=quant_config
)
self.w3 = ReplicatedLinear(
self.hidden_dim, self.ffn_dim, bias=False, linear_method=linear_method
self.hidden_dim, self.ffn_dim, bias=False, quant_config=quant_config
)
# TODO: Use vllm's SiluAndMul
@@ -76,7 +77,7 @@ class MixtralMoE(nn.Module):
def __init__(
self,
config: MixtralConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.config = config
@@ -103,7 +104,7 @@ class MixtralMoE(nn.Module):
self.num_total_experts,
config.hidden_size,
config.intermediate_size,
linear_method=linear_method,
quant_config=quant_config,
)
if idx in self.expert_indicies
else None
@@ -148,7 +149,7 @@ class MixtralAttention(nn.Module):
layer_id: int = 0,
max_position: int = 4096 * 32,
rope_theta: float = 10000,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
sliding_window: Optional[int] = None,
) -> None:
super().__init__()
@@ -180,13 +181,13 @@ class MixtralAttention(nn.Module):
self.total_num_heads,
self.total_num_kv_heads,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
hidden_size,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.rotary_emb = get_rope(
self.head_dim,
@@ -222,7 +223,7 @@ class MixtralDecoderLayer(nn.Module):
self,
config: MixtralConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
@@ -236,9 +237,9 @@ class MixtralDecoderLayer(nn.Module):
layer_id=layer_id,
rope_theta=rope_theta,
sliding_window=config.sliding_window,
linear_method=linear_method,
quant_config=quant_config,
)
self.block_sparse_moe = MixtralMoE(config=config, linear_method=linear_method)
self.block_sparse_moe = MixtralMoE(config=config, quant_config=quant_config)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(
config.hidden_size, eps=config.rms_norm_eps
@@ -273,7 +274,7 @@ class MixtralModel(nn.Module):
def __init__(
self,
config: MixtralConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.padding_idx = config.pad_token_id
@@ -286,7 +287,7 @@ class MixtralModel(nn.Module):
# config.num_hidden_layers=16
self.layers = nn.ModuleList(
[
MixtralDecoderLayer(config, i, linear_method=linear_method)
MixtralDecoderLayer(config, i, quant_config=quant_config)
for i in range(config.num_hidden_layers)
]
)
@@ -317,12 +318,12 @@ class MixtralForCausalLM(nn.Module):
def __init__(
self,
config: MixtralConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
self.linear_method = linear_method
self.model = MixtralModel(config, linear_method)
self.quant_config = quant_config
self.model = MixtralModel(config, quant_config=quant_config)
self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
self.logits_processor = LogitsProcessor(config)