chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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-12 07:37:49 +08:00
committed by GitHub
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

37
python/sglang/srt/models/commandr.py
View File

@@ -20,7 +20,7 @@
# This file is based on the LLama model definition file in transformers
"""PyTorch Cohere model."""
from typing import List, Optional, Tuple
from typing import Optional, Tuple
import torch
import torch.utils.checkpoint
@@ -29,19 +29,20 @@ from torch.nn.parameter import Parameter
from transformers import PretrainedConfig
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.linear import (
LinearMethodBase,
MergedColumnParallelLinear,
QKVParallelLinear,
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 VocabParallelEmbedding
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.utils import set_weight_attrs
from vllm.model_executor.weight_utils import (
from sglang.srt.weight_utils import (
default_weight_loader,
hf_model_weights_iterator,
)
@@ -92,7 +93,7 @@ class CohereMLP(nn.Module):
def __init__(
self,
config,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.config = config
@@ -102,13 +103,13 @@ class CohereMLP(nn.Module):
self.hidden_size,
[self.intermediate_size] * 2,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.down_proj = RowParallelLinear(
self.intermediate_size,
self.hidden_size,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.act_fn = SiluAndMul()
@@ -124,7 +125,7 @@ class CohereAttention(nn.Module):
self,
config: PretrainedConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
tp_size = get_tensor_model_parallel_world_size()
@@ -159,13 +160,13 @@ class CohereAttention(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,
self.hidden_size,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.rotary_emb = get_rope(
self.head_dim,
@@ -221,16 +222,16 @@ class CohereDecoderLayer(nn.Module):
self,
config: PretrainedConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.hidden_size = config.hidden_size
self.self_attn = CohereAttention(
config, layer_id=layer_id, linear_method=linear_method
config, layer_id=layer_id, quant_config=quant_config
)
self.mlp = CohereMLP(config, linear_method=linear_method)
self.mlp = CohereMLP(config, quant_config=quant_config)
self.input_layernorm = LayerNorm(
param_shape=(config.hidden_size), eps=config.layer_norm_eps
)
@@ -261,7 +262,7 @@ class CohereModel(nn.Module):
def __init__(
self,
config: PretrainedConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.config = config
@@ -271,7 +272,7 @@ class CohereModel(nn.Module):
)
self.layers = nn.ModuleList(
[
CohereDecoderLayer(config, i, linear_method=linear_method)
CohereDecoderLayer(config, i, quant_config=quant_config)
for i in range(config.num_hidden_layers)
]
)
@@ -303,13 +304,13 @@ class CohereForCausalLM(nn.Module):
def __init__(
self,
config: PretrainedConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
self.linear_method = linear_method
self.quant_config = quant_config
self.logits_processor = LogitsProcessor(config)
self.model = CohereModel(config, linear_method)
self.model = CohereModel(config, quant_config)
@torch.no_grad()
def forward(

39
python/sglang/srt/models/dbrx.py
View File

@@ -7,26 +7,27 @@ import torch
import torch.nn as nn
from vllm.model_executor.layers.fused_moe import fused_moe
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 (
DEFAULT_VOCAB_PADDING_SIZE,
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.utils import set_weight_attrs
from vllm.model_executor.weight_utils import (
from sglang.srt.weight_utils import (
default_weight_loader,
hf_model_weights_iterator,
)
@@ -56,7 +57,7 @@ class DbrxRouter(nn.Module):
self.num_total_experts,
bias=False,
params_dtype=params_dtype,
linear_method=None,
quant_config=None,
)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -75,7 +76,7 @@ class DbrxExperts(nn.Module):
def __init__(
self,
config: DbrxConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
params_dtype: Optional[torch.dtype] = None,
):
super().__init__()
@@ -176,7 +177,7 @@ class DbrxAttention(nn.Module):
self,
config: DbrxConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.d_model = config.d_model
@@ -194,13 +195,13 @@ class DbrxAttention(nn.Module):
self.total_num_heads,
self.total_num_kv_heads,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.out_proj = RowParallelLinear(
self.d_model,
self.d_model,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.rotary_emb = get_rope(
self.head_dim,
@@ -255,11 +256,11 @@ class DbrxFusedNormAttention(nn.Module):
self,
config: DbrxConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.d_model = config.d_model
self.attn = DbrxAttention(config, layer_id, linear_method)
self.attn = DbrxAttention(config, layer_id, quant_config=quant_config)
self.norm_1 = nn.LayerNorm(self.d_model)
self.norm_2 = nn.LayerNorm(self.d_model)
@@ -287,11 +288,11 @@ class DbrxBlock(nn.Module):
self,
config: DbrxConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.norm_attn_norm = DbrxFusedNormAttention(config, layer_id, linear_method)
self.ffn = DbrxExperts(config, linear_method)
self.norm_attn_norm = DbrxFusedNormAttention(config, layer_id, quant_config=quant_config)
self.ffn = DbrxExperts(config, quant_config=quant_config)
def forward(
self,
@@ -313,7 +314,7 @@ class DbrxModel(nn.Module):
def __init__(
self,
config: DbrxConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.wte = VocabParallelEmbedding(
@@ -321,7 +322,7 @@ class DbrxModel(nn.Module):
config.d_model,
)
self.blocks = nn.ModuleList(
[DbrxBlock(config, i, linear_method) for i in range(config.n_layers)]
[DbrxBlock(config, i, quant_config=quant_config) for i in range(config.n_layers)]
)
self.norm_f = nn.LayerNorm(config.d_model, eps=1e-5)
for module in self.modules():
@@ -351,13 +352,13 @@ class DbrxForCausalLM(nn.Module):
def __init__(
self,
config: DbrxConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.config = config
self.linear_method = linear_method
self.quant_config = quant_config
self.unpadded_vocab_size = config.vocab_size
self.transformer = DbrxModel(config, linear_method)
self.transformer = DbrxModel(config, quant_config=quant_config)
self.lm_head = ParallelLMHead(
config.vocab_size,
config.d_model,

35
python/sglang/srt/models/gemma.py
View File

@@ -10,17 +10,18 @@ from vllm.config import LoRAConfig
from vllm.model_executor.layers.activation import GeluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
LinearMethodBase,
MergedColumnParallelLinear,
QKVParallelLinear,
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 VocabParallelEmbedding
from vllm.model_executor.parallel_utils.parallel_state import (
from vllm.distributed import (
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,
)
@@ -35,17 +36,17 @@ class GemmaMLP(nn.Module):
self,
hidden_size: int,
intermediate_size: int,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
hidden_size,
[intermediate_size] * 2,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.down_proj = RowParallelLinear(
intermediate_size, hidden_size, bias=False, linear_method=linear_method
intermediate_size, hidden_size, bias=False, quant_config=quant_config,
)
self.act_fn = GeluAndMul()
@@ -66,7 +67,7 @@ class GemmaAttention(nn.Module):
layer_id: int = 0,
max_position_embeddings: int = 8192,
rope_theta: float = 10000,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.hidden_size = hidden_size
@@ -96,13 +97,13 @@ class GemmaAttention(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(
@@ -139,7 +140,7 @@ class GemmaDecoderLayer(nn.Module):
self,
config: PretrainedConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
@@ -151,12 +152,12 @@ class GemmaDecoderLayer(nn.Module):
layer_id=layer_id,
max_position_embeddings=config.max_position_embeddings,
rope_theta=config.rope_theta,
linear_method=linear_method,
quant_config=quant_config,
)
self.mlp = GemmaMLP(
hidden_size=self.hidden_size,
intermediate_size=config.intermediate_size,
linear_method=linear_method,
quant_config=quant_config,
)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(
@@ -192,7 +193,7 @@ class GemmaModel(nn.Module):
def __init__(
self,
config: PretrainedConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
@@ -203,7 +204,7 @@ class GemmaModel(nn.Module):
)
self.layers = nn.ModuleList(
[
GemmaDecoderLayer(config, i, linear_method)
GemmaDecoderLayer(config, i, quant_config=quant_config)
for i in range(config.num_hidden_layers)
]
)
@@ -264,14 +265,14 @@ class GemmaForCausalLM(nn.Module):
def __init__(
self,
config: PretrainedConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
lora_config: Optional[LoRAConfig] = None,
) -> None:
del lora_config # Unused.
super().__init__()
self.config = config
self.linear_method = linear_method
self.model = GemmaModel(config, linear_method)
self.quant_config = quant_config
self.model = GemmaModel(config, quant_config=quant_config)
self.logits_processor = LogitsProcessor(config)
@torch.no_grad()

37
python/sglang/srt/models/llama2.py
View File

@@ -1,7 +1,7 @@
# Adapted from
# https://github.com/vllm-project/vllm/blob/671af2b1c0b3ed6d856d37c21a561cc429a10701/vllm/model_executor/models/llama.py#L1
"""Inference-only LLaMA model compatible with HuggingFace weights."""
from typing import Any, Dict, List, Optional, Tuple
from typing import Any, Dict, Optional, Tuple
import torch
from torch import nn
@@ -9,20 +9,21 @@ from transformers import LlamaConfig
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
LinearMethodBase,
MergedColumnParallelLinear,
QKVParallelLinear,
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.parallel_state import (
from vllm.distributed import (
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,
)
@@ -38,17 +39,17 @@ class LlamaMLP(nn.Module):
hidden_size: int,
intermediate_size: int,
hidden_act: str,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
hidden_size,
[intermediate_size] * 2,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.down_proj = RowParallelLinear(
intermediate_size, hidden_size, bias=False, linear_method=linear_method
intermediate_size, hidden_size, bias=False, quant_config=quant_config,
)
if hidden_act != "silu":
raise ValueError(
@@ -74,7 +75,7 @@ class LlamaAttention(nn.Module):
rope_theta: float = 10000,
rope_scaling: Optional[Dict[str, Any]] = None,
max_position_embeddings: int = 8192,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.hidden_size = hidden_size
@@ -105,13 +106,13 @@ class LlamaAttention(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(
@@ -148,7 +149,7 @@ class LlamaDecoderLayer(nn.Module):
self,
config: LlamaConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
@@ -163,13 +164,13 @@ class LlamaDecoderLayer(nn.Module):
rope_theta=rope_theta,
rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
linear_method=linear_method,
quant_config=quant_config,
)
self.mlp = LlamaMLP(
hidden_size=self.hidden_size,
intermediate_size=config.intermediate_size,
hidden_act=config.hidden_act,
linear_method=linear_method,
quant_config=quant_config,
)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(
@@ -205,7 +206,7 @@ class LlamaModel(nn.Module):
def __init__(
self,
config: LlamaConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
@@ -217,7 +218,7 @@ class LlamaModel(nn.Module):
)
self.layers = nn.ModuleList(
[
LlamaDecoderLayer(config, i, linear_method)
LlamaDecoderLayer(config, i, quant_config=quant_config)
for i in range(config.num_hidden_layers)
]
)
@@ -251,12 +252,12 @@ class LlamaForCausalLM(nn.Module):
def __init__(
self,
config: LlamaConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
self.linear_method = linear_method
self.model = LlamaModel(config, linear_method)
self.quant_config = quant_config
self.model = LlamaModel(config, quant_config=quant_config)
self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
self.logits_processor = LogitsProcessor(config)

11
python/sglang/srt/models/llava.py
View File

@@ -5,10 +5,11 @@ from typing import List, Optional
import numpy as np
import torch
from torch import nn
from transformers import CLIPVisionModel, LlamaConfig, LlavaConfig
from transformers import CLIPVisionModel, LlavaConfig
from transformers.models.llava.modeling_llava import LlavaMultiModalProjector
from vllm.model_executor.layers.linear import LinearMethodBase
from vllm.model_executor.weight_utils import (
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig)
from sglang.srt.weight_utils import (
default_weight_loader,
hf_model_weights_iterator,
)
@@ -27,7 +28,7 @@ class LlavaLlamaForCausalLM(nn.Module):
def __init__(
self,
config: LlavaConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
@@ -35,7 +36,7 @@ class LlavaLlamaForCausalLM(nn.Module):
self.config.vision_config.hidden_size = config.mm_hidden_size
self.config.text_config.hidden_size = config.hidden_size
self.multi_modal_projector = LlavaMultiModalProjector(config)
self.language_model = LlamaForCausalLM(config, linear_method)
self.language_model = LlamaForCausalLM(config, quant_config=quant_config)
if "unpad" in getattr(config, "mm_patch_merge_type", ""):
self.language_model.model.image_newline = nn.Parameter(
torch.empty(config.text_config.hidden_size, dtype=torch.float16)

45
python/sglang/srt/models/mixtral.py
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)

37
python/sglang/srt/models/qwen.py
View File

@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Optional, Tuple
from typing import Any, Dict, Optional
import torch
from torch import nn
@@ -6,20 +6,21 @@ from transformers import PretrainedConfig
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
LinearMethodBase,
MergedColumnParallelLinear,
QKVParallelLinear,
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.parallel_state import (
from vllm.distributed import (
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,
)
@@ -35,7 +36,7 @@ class QWenMLP(nn.Module):
hidden_size: int,
intermediate_size: int,
hidden_act: str = "silu",
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
@@ -43,14 +44,14 @@ class QWenMLP(nn.Module):
2 * [intermediate_size],
bias=False,
gather_output=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.c_proj = RowParallelLinear(
intermediate_size,
hidden_size,
bias=False,
input_is_parallel=True,
linear_method=linear_method,
quant_config=quant_config,
)
if hidden_act != "silu":
raise ValueError(
@@ -75,7 +76,7 @@ class QWenAttention(nn.Module):
layer_id: int = 0,
rope_theta: float = 10000,
rope_scaling: Optional[Dict[str, Any]] = None,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.hidden_size = hidden_size
@@ -91,14 +92,14 @@ class QWenAttention(nn.Module):
self.head_dim,
self.total_num_heads,
bias=True,
linear_method=linear_method,
quant_config=quant_config,
)
self.c_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
hidden_size,
bias=False,
input_is_parallel=True,
linear_method=linear_method,
quant_config=quant_config,
)
self.rotary_emb = get_rope(
self.head_dim,
@@ -131,7 +132,7 @@ class QWenAttention(nn.Module):
class QWenBlock(nn.Module):
def __init__(self, config: PretrainedConfig, layer_id, linear_method=None):
def __init__(self, config: PretrainedConfig, layer_id, quant_config: Optional[QuantizationConfig] = None,):
super().__init__()
self.ln_1 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
@@ -144,7 +145,7 @@ class QWenBlock(nn.Module):
rope_theta=rope_theta,
rope_scaling=rope_scaling,
layer_id=layer_id,
linear_method=linear_method,
quant_config=quant_config,
)
self.ln_2 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
@@ -152,7 +153,7 @@ class QWenBlock(nn.Module):
self.mlp = QWenMLP(
config.hidden_size,
config.intermediate_size // 2,
linear_method=linear_method,
quant_config=quant_config,
)
def forward(
@@ -180,7 +181,7 @@ class QWenBlock(nn.Module):
class QWenModel(nn.Module):
def __init__(self, config: PretrainedConfig, linear_method=None):
def __init__(self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None,):
super().__init__()
self.config = config
self.vocab_size = config.vocab_size
@@ -192,7 +193,7 @@ class QWenModel(nn.Module):
)
self.h = nn.ModuleList(
[
QWenBlock(config, i, linear_method=linear_method)
QWenBlock(config, i, quant_config=quant_config)
for i in range(config.num_hidden_layers)
]
)
@@ -217,10 +218,10 @@ class QWenModel(nn.Module):
class QWenLMHeadModel(nn.Module):
def __init__(self, config: PretrainedConfig, linear_method=None):
def __init__(self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None,):
super().__init__()
self.config = config
self.transformer = QWenModel(config, linear_method=linear_method)
self.transformer = QWenModel(config, quant_config=quant_config)
vocab_size = ((config.vocab_size + 63) // 64) * 64
self.lm_head = ParallelLMHead(vocab_size, config.hidden_size)
self.logits_processor = LogitsProcessor(config)
@@ -275,4 +276,4 @@ class QWenLMHeadModel(nn.Module):
weight_loader(param, loaded_weight)
EntryClass = QWenLMHeadModel
EntryClass = QWenLMHeadModel

37
python/sglang/srt/models/qwen2.py
View File

@@ -1,27 +1,28 @@
# Adapted from llama2.py
# Modify details for the adaptation of Qwen2 model.
"""Inference-only Qwen2 model compatible with HuggingFace weights."""
from typing import Any, Dict, List, Optional, Tuple
from typing import Any, Dict, Optional, Tuple
import torch
from torch import nn
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
LinearMethodBase,
MergedColumnParallelLinear,
QKVParallelLinear,
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.parallel_state import (
from vllm.distributed import (
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,
)
@@ -39,17 +40,17 @@ class Qwen2MLP(nn.Module):
hidden_size: int,
intermediate_size: int,
hidden_act: str,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
hidden_size,
[intermediate_size] * 2,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.down_proj = RowParallelLinear(
intermediate_size, hidden_size, bias=False, linear_method=linear_method
intermediate_size, hidden_size, bias=False, quant_config=quant_config,
)
if hidden_act != "silu":
raise ValueError(
@@ -75,7 +76,7 @@ class Qwen2Attention(nn.Module):
rope_theta: float = 1000000,
rope_scaling: Optional[Dict[str, Any]] = None,
max_position_embeddings: int = 32768,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.hidden_size = hidden_size
@@ -106,13 +107,13 @@ class Qwen2Attention(nn.Module):
self.total_num_heads,
self.total_num_kv_heads,
bias=True,
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(
@@ -149,7 +150,7 @@ class Qwen2DecoderLayer(nn.Module):
self,
config: Qwen2Config,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
@@ -164,13 +165,13 @@ class Qwen2DecoderLayer(nn.Module):
rope_theta=rope_theta,
rope_scaling=rope_scaling,
max_position_embeddings=max_position_embeddings,
linear_method=linear_method,
quant_config=quant_config,
)
self.mlp = Qwen2MLP(
hidden_size=self.hidden_size,
intermediate_size=config.intermediate_size,
hidden_act=config.hidden_act,
linear_method=linear_method,
quant_config=quant_config,
)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(
@@ -206,7 +207,7 @@ class Qwen2Model(nn.Module):
def __init__(
self,
config: Qwen2Config,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
@@ -218,7 +219,7 @@ class Qwen2Model(nn.Module):
)
self.layers = nn.ModuleList(
[
Qwen2DecoderLayer(config, i, linear_method)
Qwen2DecoderLayer(config, i, quant_config=quant_config)
for i in range(config.num_hidden_layers)
]
)
@@ -252,12 +253,12 @@ class Qwen2ForCausalLM(nn.Module):
def __init__(
self,
config: Qwen2Config,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
self.linear_method = linear_method
self.model = Qwen2Model(config, linear_method)
self.quant_config = quant_config
self.model = Qwen2Model(config, quant_config=quant_config)
self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
self.logits_processor = LogitsProcessor(config)

31
python/sglang/srt/models/stablelm.py
View File

@@ -9,20 +9,21 @@ from torch import nn
from transformers import PretrainedConfig
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.linear import (
LinearMethodBase,
MergedColumnParallelLinear,
QKVParallelLinear,
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.parallel_state import (
from vllm.distributed import (
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,
)
@@ -34,7 +35,7 @@ from sglang.srt.managers.router.model_runner import InputMetadata
class StablelmMLP(nn.Module):
def __init__(
self, config: PretrainedConfig, linear_method: Optional[LinearMethodBase] = None
self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
@@ -44,10 +45,10 @@ class StablelmMLP(nn.Module):
config.hidden_size,
[config.intermediate_size] * 2,
bias=False,
linear_method=linear_method,
quant_config=quant_config,
)
self.down_proj = RowParallelLinear(
config.intermediate_size, config.hidden_size, bias=False
config.intermediate_size, config.hidden_size, bias=False, quant_config=quant_config,
)
self.act_fn = SiluAndMul()
@@ -63,7 +64,7 @@ class StablelmAttention(nn.Module):
self,
config: PretrainedConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
@@ -105,13 +106,11 @@ class StablelmAttention(nn.Module):
self.total_num_heads,
self.total_num_key_value_heads,
self.qkv_bias,
linear_method=linear_method,
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
self.hidden_size,
bias=False,
linear_method=linear_method,
)
self.rotary_emb = get_rope(
self.head_dim,
@@ -146,11 +145,11 @@ class StablelmDecoderLayer(nn.Module):
self,
config: PretrainedConfig,
layer_id: int = 0,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.self_attn = StablelmAttention(config, layer_id=layer_id)
self.mlp = StablelmMLP(config, linear_method)
self.mlp = StablelmMLP(config, quant_config=quant_config)
norm_eps = getattr(config, "norm_eps", getattr(config, "layer_norm_eps", 1e-05))
self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=norm_eps)
self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=norm_eps)
@@ -182,7 +181,7 @@ class StablelmDecoderLayer(nn.Module):
class StableLMEpochModel(nn.Module):
def __init__(
self, config: PretrainedConfig, linear_method: Optional[LinearMethodBase] = None
self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.embed_tokens = VocabParallelEmbedding(
@@ -191,7 +190,7 @@ class StableLMEpochModel(nn.Module):
)
self.layers = nn.ModuleList(
[
StablelmDecoderLayer(config, i, linear_method)
StablelmDecoderLayer(config, i, quant_config=quant_config)
for i in range(config.num_hidden_layers)
]
)
@@ -224,12 +223,12 @@ class StableLmForCausalLM(nn.Module):
def __init__(
self,
config: PretrainedConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
) -> None:
super().__init__()
self.config = config
self.linear_method = linear_method
self.model = StableLMEpochModel(config, linear_method)
self.quant_config = quant_config
self.model = StableLMEpochModel(config, quant_config=quant_config)
self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
self.logits_processor = LogitsProcessor(config)

2
python/sglang/srt/models/yivl.py
View File

@@ -6,7 +6,7 @@ from typing import List, Optional
import torch
import torch.nn as nn
from transformers import CLIPVisionModel, LlavaConfig
from vllm.model_executor.weight_utils import (
from sglang.srt.weight_utils import (
default_weight_loader,
hf_model_weights_iterator,
)