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 15, 09:45 UTC
[V1][Spec Decode] Share input embedding of target model with EAGLE draft model to free ~1GB for llama 3 model (#17326)
Co-authored-by: root <root@ekagra-8xh100.us-east5-a.c.serving-efficiency-poc.internal> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Ekagra Ranjan committed
May 14, 2025 at 15:31 UTC
418d2f8bfb5593bce89641d79849900f7294b859
4 files changed
+59
-19
examples/offline_inference/eagle.py
+7
@@ -105,6 +105,13 @@ def main():
105
outputs = llm.generate(prompt_token_ids=prompt_ids,
106
sampling_params=sampling_params)
107
108
+ # print the generated text
109
+ for output in outputs:
110
+ print("-" * 50)
111
+ print(f"prompt: {output.prompt}")
112
+ print(f"generated text: {output.outputs[0].text}")
113
+ print("-" * 50)
114
+
115
if not hasattr(outputs, "metrics") or outputs.metrics is None:
116
return
117
vllm/model_executor/models/llama_eagle.py
+18
-9
@@ -8,6 +8,7 @@ from transformers import LlamaConfig
8
9
from vllm.compilation.decorators import support_torch_compile
10
from vllm.config import VllmConfig
11
+from vllm.distributed.parallel_state import get_pp_group
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.logits_processor import LogitsProcessor
14
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -52,11 +53,15 @@ class LlamaModel(nn.Module):
53
self.config = vllm_config. \
54
speculative_config.draft_model_config.hf_config
55
self.vocab_size = self.config.vocab_size
55
- self.embed_tokens = VocabParallelEmbedding(
56
- self.config.vocab_size,
57
- self.config.hidden_size,
58
- prefix=maybe_prefix(prefix, "embed_tokens"),
59
- )
56
+
57
+ # if PP disabled then draft will share embed with target
58
+ if get_pp_group().world_size > 1:
59
+ self.embed_tokens = VocabParallelEmbedding(
60
+ self.config.vocab_size,
61
+ self.config.hidden_size,
62
+ prefix=maybe_prefix(prefix, "embed_tokens"),
63
+ )
64
+
65
self.layers = nn.ModuleList([
66
LlamaDecoderLayer(
67
self.config,
@@ -109,6 +114,12 @@ class LlamaModel(nn.Module):
114
weight_loader(param, loaded_weight, shard_id)
115
break
116
else:
117
+
118
+ # if PP disabled then draft will share embed with target
119
+ if get_pp_group().world_size == 1 and \
120
+ "embed_tokens." in name:
121
+ continue
122
+
123
param = params_dict[name]
124
weight_loader = getattr(param, "weight_loader",
125
default_weight_loader)
@@ -142,8 +153,7 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
153
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
154
loader = AutoWeightsLoader(
155
self,
145
- skip_prefixes=(["lm_head."]
146
- if self.config.tie_word_embeddings else None),
156
+ skip_prefixes=None,
157
)
158
159
model_weights = {}
@@ -151,5 +161,4 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
161
if "lm_head" not in name:
162
name = "model." + name
163
model_weights[name] = loaded_weight
154
-
155
- loader.load_weights(model_weights.items())
164
+ return loader.load_weights(model_weights.items())
vllm/model_executor/models/llama_eagle3.py
+10
-5
@@ -8,6 +8,7 @@ from transformers import LlamaConfig
8
9
from vllm.compilation.decorators import support_torch_compile
10
from vllm.config import VllmConfig
11
+from vllm.distributed.parallel_state import get_pp_group
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.layernorm import RMSNorm
14
from vllm.model_executor.layers.linear import QKVParallelLinear
@@ -91,11 +92,15 @@ class LlamaModel(nn.Module):
92
self.config = vllm_config. \
93
speculative_config.draft_model_config.hf_config
94
self.vocab_size = self.config.vocab_size
94
- self.embed_tokens = VocabParallelEmbedding(
95
- self.config.vocab_size,
96
- self.config.hidden_size,
97
- prefix=maybe_prefix(prefix, "embed_tokens"),
98
- )
95
+
96
+ # if PP disabled then draft will share embed with target
97
+ if get_pp_group().world_size > 1:
98
+ self.embed_tokens = VocabParallelEmbedding(
99
+ self.config.vocab_size,
100
+ self.config.hidden_size,
101
+ prefix=maybe_prefix(prefix, "embed_tokens"),
102
+ )
103
+
104
self.layers = nn.ModuleList([
105
LlamaDecoderLayer(
106
self.config,
vllm/v1/spec_decode/eagle.py
+24
-5
@@ -5,6 +5,7 @@ import torch.nn as nn
5
from vllm.attention.layer import Attention
6
from vllm.config import (CompilationLevel, VllmConfig,
7
get_layers_from_vllm_config, set_current_vllm_config)
8
+from vllm.distributed.parallel_state import get_pp_group
9
from vllm.forward_context import set_forward_context
10
from vllm.logger import init_logger
11
from vllm.model_executor.model_loader import get_model_loader
@@ -306,12 +307,30 @@ class EagleProposer:
307
self.attn_layer_name = next(iter(draft_attn_layer_names))
308
loaded_weights = self.model.load_weights(
309
loader.get_all_weights(draft_model_config, self.model))
309
- if self.vllm_config.speculative_config.method == "eagle3":
310
- if "model.embed_tokens.weight" not in loaded_weights:
311
- logger.info(
312
- "Loading EAGLE embedding weights from the target model.")
313
- self.model.model.embed_tokens = target_model.model.embed_tokens
310
+
311
+ # share embed_tokens with the target model if needed
312
+ if get_pp_group().world_size == 1:
313
+ assert "model.embed_tokens.weight" not in loaded_weights, \
314
+ "For PP = 1, Eagle draft should share embed with target model"
315
+ logger.info(
316
+ "The EAGLE head shares the same vocab embedding" \
317
+ " with the target model."
318
+ )
319
+ self.model.model.embed_tokens = target_model.model.embed_tokens
320
else:
321
+ assert "model.embed_tokens.weight" in loaded_weights, \
322
+ "For PP > 1, Eagle draft checkpoint should its own copy of "
323
+ " the model.embed_tokens.weight"
324
+ logger.info(
325
+ "Since PP > 1, the EAGLE head loaded its own vocab embedding" \
326
+ " weights instead of sharing them with the target model."
327
+ )
328
+
329
+ # share lm_head with the target model if needed
330
+ # some model definition do not define lm_head explicitly
331
+ # and reuse embed_tokens for lm_head, e.g., CohereForCausalLM
332
+ if self.vllm_config.speculative_config.method != "eagle3" and \
333
+ hasattr(target_model, "lm_head"):
334
logger.info("Loading EAGLE LM head weights from the target model.")
335
self.model.lm_head = target_model.lm_head
336