biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Misc][Refactor] Generalize linear_method to be quant_method (#4373)

Browse Source
This commit is contained in:
Cody Yu
2024-04-26 13:41:14 -07:00
committed by GitHub
parent 603ad84815
commit a62aaf1df5
45 changed files with 759 additions and 713 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
vllm/model_executor/models/jais.py
View File

@@ -29,10 +29,11 @@ from vllm.attention import Attention, AttentionMetadata
from vllm.distributed import (get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size)
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
LinearMethodBase,
QKVParallelLinear,
RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig)
from vllm.model_executor.layers.sampler import Sampler
from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding)
@@ -68,7 +69,7 @@ class JAISAttention(nn.Module):
def __init__(
self,
config: JAISConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.hidden_size = config.hidden_size
@@ -88,13 +89,13 @@ class JAISAttention(nn.Module):
self.head_dim,
total_num_heads,
bias=True,
linear_method=linear_method,
quant_config=quant_config,
)
self.c_proj = RowParallelLinear(
self.hidden_size,
self.hidden_size,
bias=True,
linear_method=linear_method,
quant_config=quant_config,
)
tp_rank = get_tensor_model_parallel_rank()
@@ -128,7 +129,7 @@ class JAISMLP(nn.Module):
self,
intermediate_size: int,
config: JAISConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
hidden_size = config.hidden_size
@@ -137,19 +138,19 @@ class JAISMLP(nn.Module):
hidden_size,
intermediate_size,
bias=True,
linear_method=linear_method,
quant_config=quant_config,
)
self.c_fc2 = (ColumnParallelLinear(
hidden_size,
intermediate_size,
bias=True,
linear_method=linear_method,
quant_config=quant_config,
) if self.swiglu else None)
self.c_proj = RowParallelLinear(
intermediate_size,
hidden_size,
bias=True,
linear_method=linear_method,
quant_config=quant_config,
)
self.act = SwiGLUActivation()
@@ -169,7 +170,7 @@ class JAISBlock(nn.Module):
def __init__(
self,
config: JAISConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
hidden_size = config.hidden_size
@@ -177,9 +178,9 @@ class JAISBlock(nn.Module):
hidden_size)
self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
self.attn = JAISAttention(config, linear_method)
self.attn = JAISAttention(config, quant_config)
self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
self.mlp = JAISMLP(inner_dim, config, linear_method)
self.mlp = JAISMLP(inner_dim, config, quant_config)
def forward(
self,
@@ -210,7 +211,7 @@ class JAISModel(nn.Module):
def __init__(
self,
config: JAISConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.config = config
@@ -227,7 +228,7 @@ class JAISModel(nn.Module):
else:
self.embeddings_scale = config.mup_embeddings_scale
self.h = nn.ModuleList([
JAISBlock(config, linear_method)
JAISBlock(config, quant_config)
for _ in range(config.num_hidden_layers)
])
self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
@@ -261,12 +262,12 @@ class JAISLMHeadModel(nn.Module):
def __init__(
self,
config: JAISConfig,
linear_method: Optional[LinearMethodBase] = None,
quant_config: Optional[QuantizationConfig] = None,
):
super().__init__()
self.config = config
self.linear_method = linear_method
self.transformer = JAISModel(config, linear_method)
self.quant_config = quant_config
self.transformer = JAISModel(config, quant_config)
self.lm_head_weight = self.transformer.wte.weight
if hasattr(config, "width_scale"):
self.output_logits_scale = config.width_scale