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

[QUANT] Add GPTQModel Dynamic Quantization + lm_head Quantization (#3790)

Browse Source 
Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>
Co-authored-by: ZX-ModelCloud <zx@modelcloud.ai>
This commit is contained in:
Qubitium-ModelCloud
2025-03-05 17:11:00 +08:00
committed by GitHub
parent 583d6af71b
commit 56a724eba3
56 changed files with 1988 additions and 282 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -24,7 +24,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import make_layers
from sglang.srt.utils import add_prefix, make_layers
@torch.jit.script
@@ -70,13 +70,14 @@ class Phi3SmallMLP(nn.Module):
2 * [self.intermediate_size],
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.up_proj",
prefix=add_prefix("up_proj", prefix),
)
self.down_proj = RowParallelLinear(
self.intermediate_size,
self.hidden_size,
bias=True,
quant_config=quant_config,
prefix=add_prefix("down_proj", prefix),
)
def forward(self, x):
@@ -140,7 +141,7 @@ class Phi3SmallSelfAttention(nn.Module):
self.num_key_value_heads,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
prefix=add_prefix("qkv_proj", prefix),
)
self.dense = RowParallelLinear(
@@ -148,7 +149,7 @@ class Phi3SmallSelfAttention(nn.Module):
self.hidden_size,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.o_proj",
prefix=add_prefix("o_proj", prefix),
)
if getattr(self.config, "rope_scaling", None) is not None:
@@ -201,6 +202,7 @@ class Phi3SmallSelfAttention(nn.Module):
self.scale,
num_kv_heads=self.num_kv_heads_per_partion,
layer_id=layer_id,
prefix=add_prefix("attn", prefix),
)
def forward(
@@ -234,13 +236,21 @@ class Phi3SmallDecoderLayer(nn.Module):
config: PretrainedConfig,
layer_id: int,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
super().__init__()
self.hidden_size = config.hidden_size
self.self_attn = Phi3SmallSelfAttention(
config, layer_id, quant_config=quant_config
config,
layer_id,
quant_config=quant_config,
prefix=add_prefix("self_attn", prefix),
)
self.mlp = Phi3SmallMLP(
config,
quant_config,
prefix=add_prefix("mlp", prefix),
)
self.mlp = Phi3SmallMLP(config, quant_config)
self.input_layernorm = nn.LayerNorm(
config.hidden_size, eps=config.layer_norm_epsilon
@@ -284,15 +294,20 @@ class Phi3SmallModel(nn.Module):
self.config = config
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size, config.hidden_size
config.vocab_size,
config.hidden_size,
prefix=add_prefix("embed_tokens", prefix),
)
self.mup_embedding_multiplier = config.mup_embedding_multiplier
self.start_layer, self.end_layer, self.layers = make_layers(
config.num_hidden_layers,
lambda prefix: Phi3SmallDecoderLayer(
config, int(prefix.split(".")[-1]), quant_config
config,
int(prefix.split(".")[-1]),
quant_config,
prefix=prefix,
),
prefix=f"{prefix}.layers",
prefix=add_prefix("layers", prefix),
)
self.final_layernorm = nn.LayerNorm(
@@ -335,6 +350,7 @@ class Phi3SmallForCausalLM(nn.Module):
self,
config: Phi3Config,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
super().__init__()
@@ -344,7 +360,7 @@ class Phi3SmallForCausalLM(nn.Module):
self.model = Phi3SmallModel(
config=config,
quant_config=quant_config,
prefix="model",
prefix=add_prefix("model", prefix),
)
self.vocab_size = config.vocab_size
self.mup_width_multiplier = config.mup_width_multiplier
@@ -354,6 +370,7 @@ class Phi3SmallForCausalLM(nn.Module):
org_num_embeddings=config.vocab_size,
padding_size=DEFAULT_VOCAB_PADDING_SIZE,
quant_config=quant_config,
prefix=add_prefix("lm_head", prefix),
)
if self.config.tie_word_embeddings:
self.lm_head.weight = self.model.embed_tokens.weight