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, 18:30 UTC
Enforce that TP > 1 is not supported for Mamba2 if Quantization is Enabled. (#14617)
Signed-off-by: Yu Chin Fabian Lim <flim@sg.ibm.com>
Yu Chin Fabian Lim committed
Mar 21, 2025 at 08:44 UTC
06dd08256f076689945418cd61397c1759f4abfa
1 file changed
+21
-16
vllm/model_executor/layers/mamba/mamba_mixer2.py
+21
-16
@@ -251,6 +251,9 @@ class MambaMixer2(CustomOp):
251
"then num_groups must equal 1."
252
)
253
254
+ assert self.tp_size == 1 or quant_config is None, \
255
+ "Tensor parallel currently not supported for quantized models."
256
+
257
self.ssm_state_size = ssm_state_size
258
self.activation = activation
259
@@ -331,22 +334,24 @@ class MambaMixer2(CustomOp):
334
], self.tp_size, tp_rank)
335
})
336
334
- delattr(self.in_proj.weight, "weight_loader")
335
- set_weight_attrs(
336
- self.in_proj.weight,
337
- {
338
- "weight_loader":
339
- mamba_v2_sharded_weight_loader(
340
- [
341
- intermediate_settings, # for gate
342
- intermediate_settings,
343
- group_shard_settings,
344
- group_shard_settings,
345
- head_setings, # for dt
346
- ],
347
- self.tp_size,
348
- tp_rank)
349
- })
337
+ if quant_config is None:
338
+ # - quant layers do not have a weight loader
339
+ delattr(self.in_proj.weight, "weight_loader")
340
+ set_weight_attrs(
341
+ self.in_proj.weight,
342
+ {
343
+ "weight_loader":
344
+ mamba_v2_sharded_weight_loader(
345
+ [
346
+ intermediate_settings, # for gate
347
+ intermediate_settings,
348
+ group_shard_settings,
349
+ group_shard_settings,
350
+ head_setings, # for dt
351
+ ],
352
+ self.tp_size,
353
+ tp_rank)
354
+ })
355
356
# - these are TPed by heads to reduce the size of the
357
# temporal shape