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 8, 10:00 UTC
[Bugfix] Fix MTP accuracy for GLM-5 (#34385)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Feb 11, 2026 at 22:08 UTC
ec12d39d44739bee408ec1473acc09e75daf1a5d
1 file changed
+18
vllm/v1/spec_decode/eagle.py
+18
@@ -1506,6 +1506,24 @@ class SpecDecodeBaseProposer:
1506
del self.model.lm_head
1507
self.model.lm_head = target_language_model.lm_head
1508
1509
+ # MTP models call compute_logits via shared_head.head (a
1510
+ # ParallelLMHead inside each MTP layer), not self.model.lm_head.
1511
+ # If the checkpoint omits a copy of the lm_head weights at the
1512
+ # MTP layer path, shared_head.head stays uninitialised and
1513
+ # produces NaN logits. Always share it explicitly.
1514
+ inner = getattr(self.model, "model", None)
1515
+ layers = getattr(inner, "layers", None) if inner else None
1516
+ if layers is not None:
1517
+ items = layers.values() if isinstance(layers, nn.ModuleDict) else layers
1518
+ for layer in items:
1519
+ sh = getattr(layer, "shared_head", None)
1520
+ if sh is not None and hasattr(sh, "head"):
1521
+ del sh.head
1522
+ sh.head = target_language_model.lm_head
1523
+ logger.info(
1524
+ "Shared target model lm_head with MTP shared_head.head."
1525
+ )
1526
+
1527
@torch.inference_mode()
1528
def dummy_run(
1529
self,