A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 13:30 UTC
[Quantization]add prefix for commandA quantized model (#17017)
Chen Xia committed
Apr 23, 2025 at 17:32 UTC
6b2427f995a81377e5758b34c8bb4c66db2f67bc
1 file changed
+8
-1
vllm/model_executor/models/commandr.py
+8
-1
@@ -89,6 +89,7 @@ class CohereMLP(nn.Module):
89
self,
90
config: CohereConfig,
91
quant_config: Optional[QuantizationConfig] = None,
92
+ prefix: str = "",
93
):
94
super().__init__()
95
self.config = config
@@ -99,12 +100,14 @@ class CohereMLP(nn.Module):
100
[self.intermediate_size] * 2,
101
bias=False,
102
quant_config=quant_config,
103
+ prefix=f"{prefix}.gate_up_proj",
104
)
105
self.down_proj = RowParallelLinear(
106
self.intermediate_size,
107
self.hidden_size,
108
bias=False,
109
quant_config=quant_config,
110
+ prefix=f"{prefix}.down_proj",
111
)
112
self.act_fn = SiluAndMul()
113
@@ -158,12 +161,14 @@ class CohereAttention(nn.Module):
161
self.total_num_kv_heads,
162
bias=False,
163
quant_config=quant_config,
164
+ prefix=f"{prefix}.qkv_proj",
165
)
166
self.o_proj = RowParallelLinear(
167
self.total_num_heads * self.head_dim,
168
self.hidden_size,
169
bias=False,
170
quant_config=quant_config,
171
+ prefix=f"{prefix}.o_proj",
172
)
173
self.rotary_emb = get_rope(
174
self.head_dim,
@@ -244,7 +249,9 @@ class CohereDecoderLayer(nn.Module):
249
quant_config=quant_config,
250
prefix=f"{prefix}.self_attn")
251
247
- self.mlp = CohereMLP(config, quant_config=quant_config)
252
+ self.mlp = CohereMLP(config,
253
+ quant_config=quant_config,
254
+ prefix=f"{prefix}.mlp")
255
self.input_layernorm = LayerNorm(param_shape=(config.hidden_size),
256
eps=config.layer_norm_eps)
257