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 4, 23:45 UTC
[Bugfix] Properly initialize `PerTensorScaleParameter` for fused-on-disk checkpoints (#39765)
Signed-off-by: Hemmi Shinichi <shemmi@preferred.jp> Signed-off-by: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Shinichi Hemmi committed
Apr 20, 2026 at 11:53 UTC
fcb31c1ac3f5af4701d45007cb2b66d084328f39
1 file changed
+18
-6
vllm/model_executor/layers/linear.py
+18
-6
@@ -916,9 +916,15 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
916
loaded_weight=loaded_weight, shard_id=idx
917
)
918
else:
919
- param.load_merged_column_weight(
920
- loaded_weight=loaded_weight, shard_id=0
921
- )
919
+ # When weights are already fused on disk (e.g. Phi-3's
920
+ # gate_up_proj), there is only a single scale for the
921
+ # entire fused matrix. Fill all slots with this scale
922
+ # to ensure that any subsequent reduction (like .max())
923
+ # works correctly while preserving the parameter shape.
924
+ for idx in range(param.data.shape[0]):
925
+ param.load_merged_column_weight(
926
+ loaded_weight=loaded_weight, shard_id=idx
927
+ )
928
return
929
elif type(param) in (RowvLLMParameter, BasevLLMParameter):
930
param.load_merged_column_weight(loaded_weight=loaded_weight)
@@ -1130,9 +1136,15 @@ class QKVParallelLinear(ColumnParallelLinear):
1136
self.validate_shard_id(loaded_shard_id)
1137
if loaded_shard_id is None: # special case for certain models
1138
if isinstance(param, PerTensorScaleParameter):
1133
- param.load_qkv_weight(
1134
- loaded_weight=loaded_weight, shard_id=0, tp_rank=self.tp_rank
1135
- )
1139
+ # When weights are already fused on disk (e.g. Phi-3's
1140
+ # qkv_proj), there is only a single scale for the entire
1141
+ # fused matrix. Fill all slots (q, k, v) with this scale
1142
+ # to ensure that any subsequent reduction (like .max())
1143
+ # works correctly while preserving the parameter shape.
1144
+ for idx in range(param.data.shape[0]):
1145
+ param.load_qkv_weight(
1146
+ loaded_weight=loaded_weight, shard_id=idx, tp_rank=self.tp_rank
1147
+ )
1148
return
1149
elif type(param) in (RowvLLMParameter, BasevLLMParameter):
1150
param.load_qkv_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank)