From a505cf807e6d3cb1a3cb71435964e031e95de5fc Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 13 May 2026 11:57:04 -0700 Subject: [PATCH] [ModelRunner V2] Share identical MTP weights (#42538) Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 90 +++++++++++-------- 1 file changed, 54 insertions(+), 36 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 39514ee7e9..fcbfc5569e 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/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