| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | import torch.nn as nn |
| 5 | |
| 6 | from vllm.config import VllmConfig, replace |
| 7 | from vllm.distributed.parallel_state import get_pp_group |
| 8 | from vllm.lora.layers.base import BaseLayerWithLoRA |
| 9 | from vllm.model_executor.model_loader import get_model |
| 10 | from vllm.model_executor.model_loader.utils import get_draft_load_config |
| 11 | from vllm.model_executor.models.utils import PPMissingLayer |
| 12 | from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config |
| 13 | |
| 14 | |
| 15 | def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: |
| 16 | """Share unless the draft declares its own copy that differs from the target. |
| 17 | |
| 18 | A draft that declares its own copy but has no top-level one (e.g. MTP |
| 19 | heads stored per layer) keeps it. |
| 20 | """ |
| 21 | if not getattr(eagle, flag, False): |
| 22 | return True |
| 23 | if draft is None or target is None: |
| 24 | return False |
| 25 | # torch.equal on GPU allocates a bool mask the size of the input. |
| 26 | # Use the faster GPU path when there is plenty of headroom; |
| 27 | # otherwise compare on CPU. |
| 28 | w = draft.weight |
| 29 | if w.is_cuda and torch.accelerator.get_memory_info(w.device)[0] < w.numel() * 2: |
| 30 | return torch.equal(w.cpu(), target.weight.cpu()) |
| 31 | return torch.equal(w, target.weight) |
| 32 | |
| 33 | |
| 34 | def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module): |
| 35 | """The target's lm_head — from get_language_model() for |
| 36 | *ForConditionalGeneration targets, else the top-level module.""" |
| 37 | return getattr(target_language_model, "lm_head", None) or getattr( |
| 38 | target_model, "lm_head", None |
| 39 | ) |
| 40 | |
| 41 | |
| 42 | def maybe_share_target_embed( |
| 43 | draft_model: nn.Module, draft_inner: nn.Module, target_inner: nn.Module |
| 44 | ) -> None: |
| 45 | """Share the target input embedding with the drafter when needed.""" |
| 46 | target_embed = getattr(target_inner, "embed_tokens", None) or getattr( |
| 47 | target_inner, "embedding", None |
| 48 | ) |
| 49 | if isinstance(target_embed, PPMissingLayer): |
| 50 | target_embed = None |
| 51 | # The drafter does not use the target's LoRA adapter. |
| 52 | if isinstance(target_embed, BaseLayerWithLoRA): |
| 53 | target_embed = target_embed.base_layer |
| 54 | draft_embed = getattr(draft_inner, "embed_tokens", None) |
| 55 | |
| 56 | if get_pp_group().world_size > 1 and not hasattr( |
| 57 | draft_model, "has_own_embed_tokens" |
| 58 | ): |
| 59 | return |
| 60 | |
| 61 | if target_embed is None: |
| 62 | if hasattr(draft_inner, "embed_tokens") and not getattr( |
| 63 | draft_model, "has_own_embed_tokens", False |
| 64 | ): |
| 65 | raise RuntimeError( |
| 66 | f"{type(draft_model).__name__} needs the target input embedding, " |
| 67 | "but it is unavailable on this PP stage" |
| 68 | ) |
| 69 | return |
| 70 | |
| 71 | if _should_share(draft_model, "has_own_embed_tokens", draft_embed, target_embed): |
| 72 | if draft_embed is not None: |
| 73 | del draft_inner.embed_tokens |
| 74 | draft_inner.embed_tokens = target_embed |
| 75 | |
| 76 | |
| 77 | def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: |
| 78 | from vllm.compilation.backends import set_model_tag |
| 79 | |
| 80 | speculative_config = vllm_config.speculative_config |
| 81 | assert speculative_config is not None |
| 82 | draft_model_config = speculative_config.draft_model_config |
| 83 | vllm_config = speculative_config.apply_draft_overrides(vllm_config) |
| 84 | draft_load_config = get_pp_safe_draft_load_config( |
| 85 | get_draft_load_config(vllm_config) |
| 86 | ) |
| 87 | if draft_load_config is not vllm_config.load_config: |
| 88 | vllm_config = replace(vllm_config, load_config=draft_load_config) |
| 89 | with set_model_tag("eagle_head"): |
| 90 | eagle_model = get_model( |
| 91 | vllm_config=vllm_config, model_config=draft_model_config |
| 92 | ) |
| 93 | |
| 94 | target_language_model = ( |
| 95 | target_model.get_language_model() |
| 96 | if hasattr(target_model, "get_language_model") |
| 97 | else target_model |
| 98 | ) |
| 99 | target_inner = target_language_model.model |
| 100 | draft_inner = eagle_model.model |
| 101 | |
| 102 | maybe_share_target_embed(eagle_model, draft_inner, target_inner) |
| 103 | |
| 104 | target_lm_head = get_target_lm_head(target_model, target_language_model) |
| 105 | draft_lm_head = getattr(eagle_model, "lm_head", None) |
| 106 | if target_lm_head is not None and _should_share( |
| 107 | eagle_model, "has_own_lm_head", draft_lm_head, target_lm_head |
| 108 | ): |
| 109 | if draft_lm_head is not None: |
| 110 | del eagle_model.lm_head |
| 111 | eagle_model.lm_head = target_lm_head |
| 112 | |
| 113 | # MTP layers route logits through layer.shared_head.head, not |
| 114 | # eagle_model.lm_head, so the per-layer copies need fixing up too. |
| 115 | layers = getattr(draft_inner, "layers", None) |
| 116 | if layers is not None: |
| 117 | items = layers.values() if isinstance(layers, nn.ModuleDict) else layers |
| 118 | for layer in items: |
| 119 | sh = getattr(layer, "shared_head", None) |
| 120 | if sh is not None and hasattr(sh, "head"): |
| 121 | del sh.head |
| 122 | sh.head = target_lm_head |
| 123 | |
| 124 | # MTP shares topk_indices_buffer with the target model. We update |
| 125 | # every module in the draft that holds a buffer reference so that |
| 126 | # the per-layer indexer and sparse-attention backends all point to |
| 127 | # the target's buffer. |
| 128 | if hasattr(target_inner, "topk_indices_buffer"): |
| 129 | target_buffer = target_inner.topk_indices_buffer |
| 130 | if target_buffer is not None: |
| 131 | for _, module in draft_inner.named_modules(): |
| 132 | if hasattr(module, "topk_indices_buffer"): |
| 133 | module.topk_indices_buffer = target_buffer |
| 134 | |
| 135 | return eagle_model |