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, 07:15 UTC
[Bugfix] Fix EAGLE3 broken logits (#18909)
Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai>
Benjamin Chislett committed
May 31, 2025 at 22:58 UTC
1bc86a3da1bd45e7d43347d6532a515950a438f0
1 file changed
+12
-11
vllm/model_executor/models/llama_eagle3.py
+12
-11
@@ -215,6 +215,9 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
215
logits = self.logits_processor(self.lm_head, hidden_states,
216
sampling_metadata)
217
if self.draft_id_to_target_id is None:
218
+ assert logits.shape[1] == self.config.vocab_size, \
219
+ "Expected logits to have shape " \
220
+ f"(*, {self.config.vocab_size}), but got {logits.shape}"
221
return logits
222
223
base = torch.arange(self.config.draft_vocab_size, device=logits.device)
@@ -234,24 +237,22 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
237
return self.model.fc(hidden_states)
238
239
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
237
- loader = AutoWeightsLoader(
238
- self,
239
- skip_prefixes=None,
240
- )
241
-
240
model_weights = {}
241
+ includes_draft_id_mapping = False
242
for name, loaded_weight in weights:
243
if "t2d" in name:
244
continue
245
if "d2t" in name:
246
name = name.replace("d2t", "draft_id_to_target_id")
247
+ includes_draft_id_mapping = True
248
elif "lm_head" not in name:
249
name = "model." + name
250
model_weights[name] = loaded_weight
251
252
- loaded_weights = loader.load_weights(model_weights.items())
253
-
254
- if 'd2t' not in loaded_weights:
255
- self.draft_id_to_target_id = None
256
-
257
- return loaded_weights
252
+ loader = AutoWeightsLoader(
253
+ self,
254
+ skip_prefixes=None,
255
+ skip_substrs=["draft_id_to_target_id"] \
256
+ if not includes_draft_id_mapping else None,
257
+ )
258
+ loader.load_weights(model_weights.items())