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 10, 14:00 UTC
[Bugfix][MTP] Fix GLM4 MoE fp8 loading with MTP on (#31757)
Signed-off-by: Andy Liu <andyliu@roblox.com>
Andy Liu committed
Jan 7, 2026 at 01:18 UTC
d111bc53ad2fbb5f28671019d21f5f753436e46d
1 file changed
+6
-1
vllm/model_executor/models/glm4_moe_mtp.py
+6
-1
@@ -106,7 +106,7 @@ class Glm4MoeMultiTokenPredictorLayer(nn.Module):
106
) -> torch.Tensor:
107
assert inputs_embeds is not None
108
# masking inputs at position 0, as not needed by MTP
109
- inputs_embeds[positions == 0] = 0
109
+ inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
110
inputs_embeds = self.enorm(inputs_embeds)
111
previous_hidden_states = self.hnorm(previous_hidden_states)
112
@@ -268,6 +268,11 @@ class Glm4MoeMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
268
if spec_layer is None:
269
continue
270
name = self._rewrite_spec_layer_name(spec_layer, name)
271
+ # Some checkpoints include weight scale tensors for the LM head even
272
+ # when the quantized head isn't built. Skip them if the model does
273
+ # not expose a matching parameter to avoid KeyError during load.
274
+ if name.endswith(".weight_scale") and name not in params_dict:
275
+ continue
276
for param_name, weight_name, shard_id in stacked_params_mapping:
277
# Skip non-stacked layers and experts (experts handled below).
278
if weight_name not in name: