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, 22:35 UTC
[Bugfix] Fix shape mismatch assertion error when loading Gemma3n model with BitsAndBytes quantization (#21808)
Signed-off-by: sydarb <areebsyed237@gmail.com>
Areeb Syed committed
Jul 30, 2025 at 09:05 UTC
fdde18229ef32d5872596ce9b004dabf310edbde
1 file changed
+24
-7
vllm/model_executor/models/gemma3n.py
+24
-7
@@ -167,22 +167,33 @@ class Gemma3nAltUp(nn.Module):
167
class Gemma3nLaurelBlock(nn.Module):
168
"""Learned Augmented Residual Layer"""
169
170
- def __init__(self, hidden_size: int, laurel_rank: int, rms_norm_eps: float,
171
- prefix: str):
170
+ def __init__(
171
+ self,
172
+ hidden_size: int,
173
+ laurel_rank: int,
174
+ rms_norm_eps: float,
175
+ *,
176
+ quant_config: Optional[QuantizationConfig] = None,
177
+ prefix: str,
178
+ ) -> None:
179
super().__init__()
180
181
self.linear_left = ColumnParallelLinear(
182
hidden_size,
183
laurel_rank,
184
bias=False,
185
+ quant_config=quant_config,
186
prefix=f"{prefix}.linear_left",
187
return_bias=False,
188
)
181
- self.linear_right = RowParallelLinear(laurel_rank,
182
- hidden_size,
183
- bias=False,
184
- prefix=f"{prefix}.linear_right",
185
- return_bias=False)
189
+ self.linear_right = RowParallelLinear(
190
+ laurel_rank,
191
+ hidden_size,
192
+ bias=False,
193
+ quant_config=quant_config,
194
+ prefix=f"{prefix}.linear_right",
195
+ return_bias=False,
196
+ )
197
self.post_laurel_norm = RMSNorm(
198
hidden_size=hidden_size,
199
eps=rms_norm_eps,
@@ -417,6 +428,7 @@ class Gemma3nDecoderLayer(nn.Module):
428
hidden_size=config.hidden_size,
429
laurel_rank=config.laurel_rank,
430
rms_norm_eps=config.rms_norm_eps,
431
+ quant_config=quant_config,
432
prefix=f"{prefix}.laurel",
433
)
434
@@ -427,6 +439,7 @@ class Gemma3nDecoderLayer(nn.Module):
439
config.hidden_size,
440
config.hidden_size_per_layer_input,
441
bias=False,
442
+ quant_config=quant_config,
443
prefix=f"{prefix}.per_layer_input_gate",
444
return_bias=False,
445
)
@@ -434,6 +447,7 @@ class Gemma3nDecoderLayer(nn.Module):
447
config.hidden_size_per_layer_input,
448
config.hidden_size,
449
bias=False,
450
+ quant_config=quant_config,
451
prefix=f"{prefix}.per_layer_projection",
452
return_bias=False,
453
)
@@ -547,6 +561,7 @@ class Gemma3nTextModel(nn.Module):
561
bias=False,
562
gather_output=True,
563
return_bias=False,
564
+ quant_config=quant_config,
565
prefix=f"{prefix}.per_layer_model_projection",
566
)
567
self.per_layer_projection_norm = RMSNorm(
@@ -566,6 +581,7 @@ class Gemma3nTextModel(nn.Module):
581
bias=False,
582
gather_output=True,
583
return_bias=False,
584
+ quant_config=quant_config,
585
prefix=f"{prefix}.{idx-1}.altup_projections",
586
) for idx in range(1, self.config.altup_num_inputs)
587
])
@@ -576,6 +592,7 @@ class Gemma3nTextModel(nn.Module):
592
bias=False,
593
gather_output=True,
594
return_bias=False,
595
+ quant_config=quant_config,
596
prefix=f"{prefix}.{idx-1}.altup_unembed_projections",
597
) for idx in range(1, self.config.altup_num_inputs)
598
])