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, 16:30 UTC
[Bugfix] [Qwen3.5]Fix Qwen3.5 FP8 quantization: tuple shard_id weight loading (#35289)
Signed-off-by: daowu.hzy <daowu.hzy@alibaba-inc.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
HZY committed
Feb 26, 2026 at 18:26 UTC
32693db8cea5cb9099c4e9d9876def97fdbc5387
1 file changed
+19
-8
vllm/model_executor/layers/linear.py
+19
-8
@@ -731,16 +731,16 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
731
loaded_shard_id: tuple[int, ...] | int | None = None,
732
):
733
self.validate_shard_id(loaded_shard_id)
734
- # FIXME(Isotr0py): Enable tuple shard_id for BNB quantization.
735
- if isinstance(loaded_shard_id, tuple):
736
- raise NotImplementedError(
737
- "Shard id with multiple indices is not supported in weight_loader, "
738
- "please use weight_loader_v2 instead."
739
- )
734
# Special case for GGUF
735
# initialize GGUF param after we know the quantize type
736
is_gguf_weight = getattr(param, "is_gguf_weight", False)
737
is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False)
738
+ if isinstance(loaded_shard_id, tuple) and (
739
+ is_gguf_weight or is_gguf_weight_type
740
+ ):
741
+ raise NotImplementedError(
742
+ "Shard id with multiple indices is not supported for GGUF."
743
+ )
744
if is_gguf_weight_type:
745
if loaded_shard_id is not None:
746
param.data[loaded_shard_id].copy_(loaded_weight)
@@ -768,7 +768,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
768
# Special case for per-tensor scale to load scalar into fused array.
769
needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
770
771
- if loaded_shard_id is None:
771
+ if loaded_shard_id is None or isinstance(loaded_shard_id, tuple):
772
# Loaded weight is already fused on disk (mlp).
773
# (e.g., Phi-3's gate_up_proj).
774
if output_dim is None:
@@ -780,10 +780,21 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
780
assert param_data.shape == loaded_weight.shape
781
param_data.copy_(loaded_weight)
782
return
783
+
784
+ output_sizes = (
785
+ self.output_sizes[loaded_shard_id[0] : loaded_shard_id[-1] + 1]
786
+ if loaded_shard_id is not None
787
+ else self.output_sizes
788
+ )
789
current_shard_offset = 0
790
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
791
+ if use_bitsandbytes_4bit and isinstance(loaded_shard_id, tuple):
792
+ raise NotImplementedError(
793
+ "Shard id with multiple indices is not supported "
794
+ "for BNB quantization yet."
795
+ )
796
shard_offsets: list[tuple[int, int, int]] = []
786
- for i, output_size in enumerate(self.output_sizes):
797
+ for i, output_size in enumerate(output_sizes):
798
shard_offsets.append((i, current_shard_offset, output_size))
799
current_shard_offset += output_size
800
packed_dim = getattr(param, "packed_dim", None)