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 2, 08:30 UTC
fix: glm5.1 pp model loading (#42944)
Signed-off-by: UranusSeven <109661872+UranusSeven@users.noreply.github.com>
Uranus committed
Jun 1, 2026 at 15:14 UTC
1f6048abe57511ed789deb8f9db4760546bcf4f5
2 files changed
+25
-5
vllm/model_executor/models/deepseek_mtp.py
+8
-2
@@ -35,7 +35,7 @@ from .deepseek_v2 import (
35
_try_load_fp8_indexer_wk,
36
get_spec_layer_idx_from_weight_name,
37
)
38
-from .utils import maybe_prefix
38
+from .utils import get_pp_missing_layer_names, maybe_prefix
39
40
logger = init_logger(__name__)
41
@@ -267,6 +267,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
267
),
268
)
269
270
+ pp_missing_layer_names = get_pp_missing_layer_names(self)
271
params_dict = dict(self.named_parameters())
272
loaded_params: set[str] = set()
273
_pending_wk_fp8: dict = {} # FP8 indexer wk dequant buffer
@@ -282,7 +283,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
283
name = self._rewrite_spec_layer_name(spec_layer, name)
284
285
if _try_load_fp8_indexer_wk(
285
- name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
286
+ name,
287
+ loaded_weight,
288
+ _pending_wk_fp8,
289
+ params_dict,
290
+ loaded_params,
291
+ pp_missing_layer_names,
292
):
293
continue
294
vllm/model_executor/models/deepseek_v2.py
+17
-3
@@ -105,6 +105,7 @@ from .interfaces import (
105
)
106
from .utils import (
107
PPMissingLayer,
108
+ get_pp_missing_layer_names,
109
is_pp_missing_parameter,
110
make_empty_intermediate_tensors_factory,
111
make_layers,
@@ -742,7 +743,9 @@ class Indexer(nn.Module):
743
return self.indexer_op(hidden_states, q_fp8, k, weights)
744
745
745
-def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
746
+def _try_load_fp8_indexer_wk(
747
+ name, tensor, buf, params_dict, loaded_params, pp_missing_layer_names
748
+):
749
"""
750
We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
751
in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
@@ -758,6 +761,12 @@ def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
761
return False # WK is not in FP8 format, ignore.
762
# Buffer this tensor (weight or scale) until both have arrived.
763
layer_prefix = name.rsplit(".wk.", 1)[0] # e.g. "model.layers.0.self_attn.indexer"
764
+ fused_name = f"{layer_prefix}.wk_weights_proj.weight"
765
+ if any(
766
+ name.startswith(missing_layer_name)
767
+ for missing_layer_name in pp_missing_layer_names
768
+ ):
769
+ return True
770
entry = buf.setdefault(layer_prefix, {})
771
entry["weight" if is_weight else "scale"] = tensor
772
if "weight" not in entry or "scale" not in entry:
@@ -775,7 +784,6 @@ def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
784
)
785
786
# Load the dequantized weight into shard 0 of the fused buffer.
778
- fused_name = f"{layer_prefix}.wk_weights_proj.weight"
787
param = params_dict[fused_name]
788
param.weight_loader(param, weight_bf16, 0)
789
loaded_params.add(fused_name)
@@ -1379,6 +1387,7 @@ class DeepseekV2Model(nn.Module):
1387
num_redundant_experts=self.num_redundant_experts,
1388
)
1389
1390
+ pp_missing_layer_names = get_pp_missing_layer_names(self)
1391
params_dict = dict(self.named_parameters())
1392
loaded_params: set[str] = set()
1393
for name, loaded_weight in weights:
@@ -1394,7 +1403,12 @@ class DeepseekV2Model(nn.Module):
1403
)
1404
1405
if _try_load_fp8_indexer_wk(
1397
- name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
1406
+ name,
1407
+ loaded_weight,
1408
+ _pending_wk_fp8,
1409
+ params_dict,
1410
+ loaded_params,
1411
+ pp_missing_layer_names,
1412
):
1413
continue
1414