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, 17:15 UTC
[Feature] Enhance EAGLE Architecture with Proper RMS Norms (#14990)
Signed-off-by: Bryan Lu <yuzhelu@amazon.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Bryan Lu committed
Mar 26, 2025 at 01:24 UTC
781d0562809b34f0c548cd354bbc01c861814f94
3 files changed
+70
-12
vllm/config.py
+12
-4
@@ -800,10 +800,18 @@ class ModelConfig:
800
801
@property
802
def is_deepseek_mla(self) -> bool:
803
- return (hasattr(self.hf_text_config, "model_type")) \
804
- and (self.hf_text_config.model_type in \
805
- ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp'))\
806
- and (self.hf_text_config.kv_lora_rank is not None)
803
+ if not hasattr(self.hf_text_config, "model_type"):
804
+ return False
805
+ elif self.hf_text_config.model_type in \
806
+ ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp'):
807
+ return self.hf_text_config.kv_lora_rank is not None
808
+ elif self.hf_text_config.model_type == 'eagle':
809
+ # if the model is an EAGLE module, check for the
810
+ # underlying architecture
811
+ return self.hf_text_config.model.model_type in \
812
+ ('deepseek_v2', 'deepseek_v3') \
813
+ and self.hf_text_config.kv_lora_rank is not None
814
+ return False
815
816
def get_head_size(self) -> int:
817
# TODO remove hard code
vllm/model_executor/models/eagle.py
+45
-6
@@ -7,6 +7,7 @@ import torch.nn as nn
7
8
from vllm.config import VllmConfig
9
from vllm.logger import init_logger
10
+from vllm.model_executor.layers.layernorm import RMSNorm
11
from vllm.model_executor.layers.logits_processor import LogitsProcessor
12
from vllm.model_executor.layers.sampler import SamplerOutput
13
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -59,7 +60,15 @@ class EAGLE(nn.Module):
60
truncated_vocab_size < vocab_size. To use this technique, one has to find
61
the top-k most frequent tokens in target dataset and add that as a tensor
62
in the draft checkpoint (using key token_map). Also, the draft config
62
- needs to have truncated_vocab_size (=k) as an attribute."""
63
+ needs to have truncated_vocab_size (=k) as an attribute.
64
+ 4. We allow an enhanced EAGLE architecture similar to the DeepSeek MTP
65
+ module with regards to the use of additional RMS norms. The original
66
+ EAGLE architecture 1) skips the pre-attention norm in its first
67
+ transformer block, and 2) skips the final output norm, both of which we
68
+ found to be suboptimal. We also add the support for separate norms
69
+ applying to both the token embedding and hidden states before projection
70
+ as in DeepSeek MTP, which we found to improve performance as well.
71
+ """
72
73
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
74
super().__init__()
@@ -81,9 +90,22 @@ class EAGLE(nn.Module):
90
# While weights and biases are generally not needed,
91
# they are retained here to support certain unit tests
92
# (e.g., spec_decode/e2e/test_eagle_correctness.py).
84
- self.model.model.layers[0].input_layernorm = DummyInputLayerNorm(
85
- weight=self.model.model.layers[0].input_layernorm.weight)
86
- self.model.model.norm = DummyOutputNorm()
93
+ if not hasattr(self.config.model,
94
+ "skip_prenorm") or self.config.model.skip_prenorm:
95
+ self.model.model.layers[0].input_layernorm = DummyInputLayerNorm(
96
+ weight=self.model.model.layers[0].input_layernorm.weight)
97
+
98
+ if not hasattr(
99
+ self.config.model,
100
+ "skip_output_norm") or self.config.model.skip_output_norm:
101
+ self.model.model.norm = DummyOutputNorm()
102
+
103
+ self.add_para_norm = False
104
+ if hasattr(self.config.model,
105
+ "add_para_norm") and self.config.model.add_para_norm:
106
+ self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
107
+ self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
108
+ self.add_para_norm = True
109
110
self.orig_vocab_size = config.vocab_size
111
self.truncated_vocab_size = config.truncated_vocab_size
@@ -128,8 +150,17 @@ class EAGLE(nn.Module):
150
if inputs_embeds is None:
151
inputs_embeds = self.get_input_embeddings(input_ids)
152
131
- inputs_embeds = self.fc(
132
- torch.cat([inputs_embeds, previous_hidden_states], dim=-1))
153
+ if self.add_para_norm:
154
+ inputs_embeds = torch.cat([
155
+ self.enorm(inputs_embeds),
156
+ self.hnorm(previous_hidden_states)
157
+ ],
158
+ dim=-1)
159
+ else:
160
+ inputs_embeds = torch.cat([inputs_embeds, previous_hidden_states],
161
+ dim=-1)
162
+
163
+ inputs_embeds = self.fc(inputs_embeds)
164
165
inputs_embeds[positions == 0] = 0 # masking inputs at position=0
166
@@ -190,6 +221,14 @@ class EAGLE(nn.Module):
221
else:
222
logger.warning_once("Found bias in the loaded weights but "
223
"the model config doesn't have bias.")
224
+ elif name.startswith("enorm.weight"):
225
+ weight_loader = getattr(self.enorm.weight, "weight_loader",
226
+ default_weight_loader)
227
+ weight_loader(self.enorm.weight, loaded_weight)
228
+ elif name.startswith("hnorm.weight"):
229
+ weight_loader = getattr(self.hnorm.weight, "weight_loader",
230
+ default_weight_loader)
231
+ weight_loader(self.hnorm.weight, loaded_weight)
232
elif name.startswith("model.lm_head.") or name.startswith(
233
"model.model."):
234
model_weights[name.split("model.", 1)[-1]] = loaded_weight
vllm/transformers_utils/configs/eagle.py
+13
-2
@@ -5,6 +5,8 @@ from typing import Optional, Union
5
6
from transformers import AutoConfig, PretrainedConfig
7
8
+from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekV2Config
9
+
10
11
class EAGLEConfig(PretrainedConfig):
12
model_type = "eagle"
@@ -14,8 +16,17 @@ class EAGLEConfig(PretrainedConfig):
16
truncated_vocab_size: Optional[int] = None,
17
**kwargs):
18
17
- model_config = None if model is None else (AutoConfig.for_model(
18
- **model) if isinstance(model, dict) else model)
19
+ model_config: Union[PretrainedConfig, DeepseekV2Config, None]
20
+ if isinstance(model, dict):
21
+ archs = model.get("architectures", [])
22
+ target_archs = ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]
23
+ if any(target_arch in archs for target_arch in target_archs):
24
+ # AutoConfig does not support DeepSeek MoE models yet
25
+ model_config = DeepseekV2Config(**model)
26
+ else:
27
+ model_config = AutoConfig.for_model(**model)
28
+ else:
29
+ model_config = model
30
31
for k, v in kwargs.items():
32
if k != "architectures" and k != "model_type" and hasattr(