[ModelRunner V2] Share identical MTP weights (#42538)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

vllm/v1/worker/gpu/spec_decode/eagle/utils.py

@@ -1,11 +1,29 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
 import torch.nn as nn
 
 from vllm.config import VllmConfig
+from vllm.distributed.parallel_state import get_pp_group
 from vllm.model_executor.model_loader import get_model
 
 
+def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
+    """Share when the draft has no own copy, or its copy matches the target."""
+
+    if not getattr(eagle, flag, False) or draft is None:
+        return True
+    if target is None:
+        return False
+    # torch.equal on GPU allocates a bool mask the size of the input.
+    # Use the faster GPU path when there is plenty of headroom;
+    # otherwise compare on CPU.
+    w = draft.weight
+    if w.is_cuda and torch.cuda.mem_get_info(w.device)[0] < w.numel() * 2:
+        return torch.equal(w.cpu(), target.weight.cpu())
+    return torch.equal(w, target.weight)
+
+
 def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
     from vllm.compilation.backends import set_model_tag
 
@@ -17,51 +35,51 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
             vllm_config=vllm_config, model_config=draft_model_config
         )
 
-    # Share target embeddings when the draft checkpoint does not include
-    # its own vocab embedding table.
-    share_embeddings = True
-    if hasattr(eagle_model, "has_own_embed_tokens"):
-        share_embeddings = not eagle_model.has_own_embed_tokens
-    if share_embeddings:
-        target_language_model = (
-            target_model.get_language_model()
-            if hasattr(target_model, "get_language_model")
-            else target_model
+    target_language_model = (
+        target_model.get_language_model()
+        if hasattr(target_model, "get_language_model")
+        else target_model
+    )
+    target_inner = target_language_model.model
+    draft_inner = eagle_model.model
+
+    # Skip embedding sharing under PP — each rank owns its own embedding.
+    if get_pp_group().world_size == 1:
+        target_embed = getattr(target_inner, "embed_tokens", None) or getattr(
+            target_inner, "embedding", None
         )
-        inner_model = getattr(target_language_model, "model", None)
-        target_embed_tokens = None
-        if inner_model is not None:
-            if hasattr(inner_model, "embed_tokens"):
-                target_embed_tokens = inner_model.embed_tokens
-            elif hasattr(inner_model, "embedding"):
-                target_embed_tokens = inner_model.embedding
-        if target_embed_tokens is not None and hasattr(eagle_model, "model"):
-            if hasattr(eagle_model.model, "embed_tokens"):
-                del eagle_model.model.embed_tokens
-            eagle_model.model.embed_tokens = target_embed_tokens
+        draft_embed = getattr(draft_inner, "embed_tokens", None)
+        if target_embed is not None and _should_share(
+            eagle_model, "has_own_embed_tokens", draft_embed, target_embed
+        ):
+            if draft_embed is not None:
+                del draft_inner.embed_tokens
+            draft_inner.embed_tokens = target_embed
 
-    # Only share target lm_head when the draft model does not own one.
-    share_lm_head = True
-    if hasattr(eagle_model, "has_own_lm_head"):
-        share_lm_head = not eagle_model.has_own_lm_head
-    if share_lm_head and hasattr(target_model, "lm_head"):
-        if hasattr(eagle_model, "lm_head"):
+    target_lm_head = getattr(target_model, "lm_head", None)
+    draft_lm_head = getattr(eagle_model, "lm_head", None)
+    if target_lm_head is not None and _should_share(
+        eagle_model, "has_own_lm_head", draft_lm_head, target_lm_head
+    ):
+        if draft_lm_head is not None:
             del eagle_model.lm_head
-        eagle_model.lm_head = target_model.lm_head
+        eagle_model.lm_head = target_lm_head
 
-        # MTP models call compute_logits via shared_head.head (a
-        # ParallelLMHead inside each MTP layer), not self.model.lm_head.
-        # If the checkpoint omits a copy of the lm_head weights at the
-        # MTP layer path, shared_head.head stays uninitialised and
-        # produces zero/NaN logits. Share it explicitly from the target.
-        inner = getattr(eagle_model, "model", None)
-        layers = getattr(inner, "layers", None) if inner is not None else None
+        # MTP layers route logits through layer.shared_head.head, not
+        # eagle_model.lm_head, so the per-layer copies need fixing up too.
+        layers = getattr(draft_inner, "layers", None)
         if layers is not None:
             items = layers.values() if isinstance(layers, nn.ModuleDict) else layers
             for layer in items:
                 sh = getattr(layer, "shared_head", None)
                 if sh is not None and hasattr(sh, "head"):
                     del sh.head
-                    sh.head = target_model.lm_head
+                    sh.head = target_lm_head
+
+    # MTP also shares a topk_indices_buffer between target and draft.
+    if hasattr(target_inner, "topk_indices_buffer"):
+        if hasattr(draft_inner, "topk_indices_buffer"):
+            del draft_inner.topk_indices_buffer
+        draft_inner.topk_indices_buffer = target_inner.topk_indices_buffer
 
     return eagle_model