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
[Bugfix] Fix BNB name match (#24735)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Sep 12, 2025 at 19:12 UTC
60a095192499bf10b588b7e34e6d3cba6861f01e
1 file changed
+6
-5
vllm/model_executor/model_loader/bitsandbytes_loader.py
+6
-5
@@ -326,7 +326,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):
326
327
global_tp_size = get_tensor_model_parallel_world_size()
328
global_tp_rank = get_tensor_model_parallel_rank()
329
-
329
+ check_match = (lambda weight_name, module_name: weight_name.
330
+ removesuffix(".weight") == module_name)
331
for (
332
org_weight_name,
333
mapped_weight_name,
@@ -347,12 +348,12 @@ class BitsAndBytesModelLoader(BaseModelLoader):
348
) and mapped_weight_name.endswith(".weight"):
349
# Without sharding
350
if any(
350
- mapped_weight_name.startswith(module)
351
+ check_match(mapped_weight_name, module)
352
for module in self.unsharded_weights_modules):
353
weight_sub_tensor = weight_tensor
354
# Shard by column
355
elif any(
355
- mapped_weight_name.startswith(module)
356
+ check_match(mapped_weight_name, module)
357
for module in self.column_sharded_weights_modules):
358
total_size = weight_tensor.size(-1)
359
start_index = total_size // tp_size * tp_rank
@@ -362,14 +363,14 @@ class BitsAndBytesModelLoader(BaseModelLoader):
363
# Weights have fused on disk. In this case, we assume that the
364
# weight and module use same name.
365
elif any(
365
- mapped_weight_name.startswith(module)
366
+ check_match(mapped_weight_name, module)
367
for module in self.maybe_fused_weights_modules):
368
# special case for fused weights
369
# get the size of each shard weight tensor
370
total_shard_sizes = next(
371
(sizes for module, sizes in
372
self.maybe_fused_weights_modules.items()
372
- if mapped_weight_name.startswith(module)))
373
+ if check_match(mapped_weight_name, module)))
374
total_size = weight_tensor.size(0)
375
assert total_size == sum(total_shard_sizes)
376
# get the start/end index of each shard weight tensor