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 14, 05:30 UTC
Fix loading of quantized BigCode models (#22463)
Signed-off-by: Eldar Kurtic <eldar@neuralmagic.com>
Eldar Kurtić committed
Aug 9, 2025 at 08:12 UTC
10a02535d4252353880486f6fdf91e5ce7507977
1 file changed
+14
-4
vllm/model_executor/models/gpt_bigcode.py
+14
-4
@@ -45,7 +45,8 @@ from vllm.sequence import IntermediateTensors
45
46
from .interfaces import SupportsLoRA, SupportsPP
47
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
48
- make_empty_intermediate_tensors_factory, make_layers)
48
+ make_empty_intermediate_tensors_factory, make_layers,
49
+ maybe_prefix)
50
51
52
class GPTBigCodeAttention(nn.Module):
@@ -83,6 +84,7 @@ class GPTBigCodeAttention(nn.Module):
84
total_num_kv_heads,
85
bias=True,
86
quant_config=quant_config,
87
+ prefix=f"{prefix}.c_attn",
88
)
89
90
self.c_proj = RowParallelLinear(
@@ -90,6 +92,7 @@ class GPTBigCodeAttention(nn.Module):
92
self.hidden_size,
93
bias=True,
94
quant_config=quant_config,
95
+ prefix=f"{prefix}.c_proj",
96
)
97
self.attn = Attention(self.num_heads,
98
self.head_dim,
@@ -123,6 +126,7 @@ class GPTBigMLP(nn.Module):
126
intermediate_size: int,
127
config: GPTBigCodeConfig,
128
quant_config: Optional[QuantizationConfig] = None,
129
+ prefix: str = "",
130
):
131
super().__init__()
132
hidden_size = config.hidden_size
@@ -131,12 +135,14 @@ class GPTBigMLP(nn.Module):
135
intermediate_size,
136
bias=True,
137
quant_config=quant_config,
138
+ prefix=f"{prefix}.c_fc",
139
)
140
self.c_proj = RowParallelLinear(
141
intermediate_size,
142
hidden_size,
143
bias=True,
144
quant_config=quant_config,
145
+ prefix=f"{prefix}.c_proj",
146
)
147
self.act = get_act_fn(config.activation_function)
148
@@ -167,7 +173,10 @@ class GPTBigCodeBlock(nn.Module):
173
quant_config,
174
prefix=f"{prefix}.attn")
175
self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
170
- self.mlp = GPTBigMLP(inner_dim, config, quant_config)
176
+ self.mlp = GPTBigMLP(inner_dim,
177
+ config,
178
+ quant_config,
179
+ prefix=f"{prefix}.mlp")
180
181
def forward(
182
self,
@@ -260,7 +269,7 @@ class GPTBigCodeModel(nn.Module):
269
weight_loader = getattr(param, "weight_loader",
270
default_weight_loader)
271
# TODO (@robertgshaw2-neuralmagic): move to fp8 linear method
263
- if "c_attn.input_scale" in name or "c_attn.weight_scale" in name:
272
+ if "c_attn.input_scale" in name:
273
weight_loader(param, loaded_weight, 'q')
274
weight_loader(param, loaded_weight, 'k')
275
weight_loader(param, loaded_weight, 'v')
@@ -284,7 +293,8 @@ class GPTBigCodeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
293
294
self.quant_config = quant_config
295
self.transformer = GPTBigCodeModel(vllm_config=vllm_config,
287
- prefix=prefix)
296
+ prefix=maybe_prefix(
297
+ prefix, "transformer"))
298
if self.config.tie_word_embeddings:
299
self.lm_head = self.transformer.wte
300
else: