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 7, 07:45 UTC
Enable bnb for multiple indices weight (#35838)
Signed-off-by: xjx <493337577@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
xjx committed
Mar 4, 2026 at 09:46 UTC
9a9d4424649fc360346bc63fd395c3f62731b7cf
1 file changed
+14
-5
vllm/model_executor/layers/linear.py
+14
-5
@@ -744,10 +744,14 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
744
)
745
current_shard_offset = 0
746
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
747
- if use_bitsandbytes_4bit and isinstance(loaded_shard_id, tuple):
747
+ if (
748
+ use_bitsandbytes_4bit
749
+ and isinstance(loaded_shard_id, tuple)
750
+ and self.tp_size > 1
751
+ ):
752
raise NotImplementedError(
753
"Shard id with multiple indices is not supported "
750
- "for BNB quantization yet."
754
+ "for BNB quantization with TP yet."
755
)
756
shard_offsets: list[tuple[int, int, int]] = []
757
for i, output_size in enumerate(output_sizes):
@@ -815,9 +819,14 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
819
is_sharded_weight = is_sharded_weight or use_bitsandbytes_4bit
820
821
if use_bitsandbytes_4bit:
818
- shard_size = loaded_weight.shape[output_dim]
819
- shard_offset = loaded_weight.shape[output_dim] * loaded_shard_id
820
-
822
+ index = list(itertools.accumulate([0] + self.output_sizes))
823
+ orig_offsets = {
824
+ str(i): (index[i], size) for i, size in enumerate(self.output_sizes)
825
+ }
826
+ orig_offsets["total"] = (self.output_size, 0)
827
+ shard_size, shard_offset = adjust_bitsandbytes_4bit_shard(
828
+ param, orig_offsets, str(loaded_shard_id)
829
+ )
830
param_data = param_data.narrow(output_dim, shard_offset, shard_size)
831
start_idx = self.tp_rank * shard_size
832
if not is_sharded_weight: