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, 16:00 UTC
[Model] Add module name prefixes to gemma3 (#15889)
Signed-off-by: Bartholomew Sabat <bartek@recursal.ai> Co-authored-by: Bartholomew Sabat <bartek@recursal.ai>
cloud11665 committed
Apr 2, 2025 at 02:13 UTC
9ec8257914b129080a8b4eab65da400eaaca6300
1 file changed
+17
-6
vllm/model_executor/models/gemma3.py
+17
-6
@@ -59,16 +59,23 @@ class Gemma3MLP(nn.Module):
59
intermediate_size: int,
60
hidden_activation: str,
61
quant_config: Optional[QuantizationConfig] = None,
62
+ prefix: str = "",
63
) -> None:
64
super().__init__()
65
self.gate_up_proj = MergedColumnParallelLinear(
65
- hidden_size, [intermediate_size] * 2,
66
+ hidden_size,
67
+ [intermediate_size] * 2,
68
+ bias=False,
69
+ quant_config=quant_config,
70
+ prefix=f"{prefix}.gate_up_proj",
71
+ )
72
+ self.down_proj = RowParallelLinear(
73
+ intermediate_size,
74
+ hidden_size,
75
bias=False,
67
- quant_config=quant_config)
68
- self.down_proj = RowParallelLinear(intermediate_size,
69
- hidden_size,
70
- bias=False,
71
- quant_config=quant_config)
76
+ quant_config=quant_config,
77
+ prefix=f"{prefix}.down_proj",
78
+ )
79
if hidden_activation != "gelu_pytorch_tanh":
80
raise ValueError(
81
"Gemma3 uses `gelu_pytorch_tanh` as the hidden activation "
@@ -125,12 +132,14 @@ class Gemma3Attention(nn.Module):
132
self.total_num_kv_heads,
133
bias=config.attention_bias,
134
quant_config=quant_config,
135
+ prefix=f"{prefix}.qkv_proj",
136
)
137
self.o_proj = RowParallelLinear(
138
self.total_num_heads * self.head_dim,
139
hidden_size,
140
bias=config.attention_bias,
141
quant_config=quant_config,
142
+ prefix=f"{prefix}.o_proj",
143
)
144
145
self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
@@ -293,6 +302,7 @@ class Gemma3DecoderLayer(nn.Module):
302
intermediate_size=config.intermediate_size,
303
hidden_activation=config.hidden_activation,
304
quant_config=quant_config,
305
+ prefix=f"{prefix}.mlp",
306
)
307
self.input_layernorm = GemmaRMSNorm(config.hidden_size,
308
eps=config.rms_norm_eps)
@@ -344,6 +354,7 @@ class Gemma3Model(nn.Module):
354
self.embed_tokens = VocabParallelEmbedding(
355
config.vocab_size,
356
config.hidden_size,
357
+ prefix=f"{prefix}.embed_tokens",
358
)
359
self.start_layer, self.end_layer, self.layers = make_layers(
360
config.num_hidden_layers,