[Model] Support deepseek with eagle #21086
Conversation
mergify
bot
added
deepseek
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for Eagle speculative decoding with Deepseek models. I've found a few critical issues in the implementation that will prevent it from working correctly. The model implementation in deepseek_eagle.py incorrectly handles hidden state dimensions and is missing the lm_head layer, which will cause runtime errors. Additionally, the model registry key in registry.py seems to be incorrect, which would prevent the model from being loaded.
| self.fc = nn.Linear( | ||
| self.config.model.hidden_size * 2, | ||
| self.config.model.hidden_size, | ||
| bias=False, | ||
| ) | ||
|
|
||
| self.enorm = RMSNorm(self.config.hidden_size, | ||
| eps=self.config.rms_norm_eps) | ||
| self.hnorm = RMSNorm(self.config.hidden_size, | ||
| eps=self.config.rms_norm_eps) | ||
| self.norm = RMSNorm(self.config.hidden_size, | ||
| eps=self.config.rms_norm_eps) |
There was a problem hiding this comment.
The implementation of DeepseekV2Model assumes that the draft model and the target model share the same hidden size. For instance, self.hnorm is initialized with the draft model's hidden size (self.config.hidden_size) but is applied to hidden_states from the target model.
This assumption is incorrect for the models used in testing (deepseek-r1 has a hidden size of 4096, while eagle-deepseek-r1 has 1024), and will lead to a runtime error due to shape mismatch.
To fix this, you should explicitly use the hidden sizes from both the draft and target model configurations. You can access the target model's configuration via vllm_config.model_config.
target_config = vllm_config.model_config.hf_config
draft_hidden_size = self.config.hidden_size
target_hidden_size = target_config.hidden_size
self.fc = nn.Linear(
draft_hidden_size + target_hidden_size,
draft_hidden_size,
bias=False,
)
self.enorm = RMSNorm(draft_hidden_size,
eps=self.config.rms_norm_eps)
self.hnorm = RMSNorm(target_hidden_size,
eps=target_config.rms_norm_eps)
self.norm = RMSNorm(draft_hidden_size,
eps=self.config.rms_norm_eps)
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
|
Hi! I tried installing this pr from source. But got Should the |
Thanks for your comment, fixed now. |
xyang16
changed the title
Amazing work! I wonder if you could share how you got eagle618/eagle-deepseek-r1? As this pr could also improve DS V3 etc. Thank you! |
xyang16
force-pushed
the
eagle
branch
3 times, most recently
from
96a0839 to
6981998
Compare
Signed-off-by: Xin Yang <xyangx@amazon.com>
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
This PR is to support running eagle speculative decoding on deepseek model. Changed the following file:
deepseek_eagle.py: deepseek eagle model definitionregistry.py: add the model to registryTest Plan
We have ran Deepseek with an eagle draft model.
Test Result
Acceptance rate is around 60% for vanilla eagle head (and better acceptance rate for fine tuned eagle head).
(Optional) Documentation Update