Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Spec Decode] Disable Log Prob serialization to CPU for spec decoding for both draft and target models. #6485

Merged
merged 45 commits into from
Jul 21, 2024
Merged

[Spec Decode] Disable Log Prob serialization to CPU for spec decoding for both draft and target models. #6485

merged 45 commits into from
Jul 21, 2024

Conversation

sroy745
Copy link
Collaborator

@sroy745 sroy745 commented Jul 16, 2024 •
edited
Loading

In this PR we disable the serialization of the LogProbs to CPU for both draft and target models. To that end we make the following changes

  1. Add top level flags to allow users to configure if they want token log probabilities during speculative decoding
  2. Introduce a light weight TargetModelRunner using which we can set the SamplingMetadata.skip_sampler_cpu_output as needed.
  3. Made changes to SpecDecodeWorker to skip serialization of log probability tensors if not needed.
  4. Fixed clearing out the appropriate tensor in SamplerOutput. E.g. earlier we were clearing sampler_output.probs which is not a valid tensor.
  5. Fixed failing tests as needed.

sroy745 and others added 19 commits May 28, 2024 20:39
@sroy745
Merge pull request #1 from vllm-project/main
5650b95 
Pull from head
@sroy745
Merge branch 'vllm-project:main' into main
8f36146
@sroy745
Merge branch 'vllm-project:main' into main
9e75057
@sroy745
Merge branch 'vllm-project:main' into main
db2c679
@sroy745
Merge branch 'vllm-project:main' into main
8d7512c
@sroy745
Merge branch 'vllm-project:main' into main
1473f74
@sroy745
Merge branch 'vllm-project:main' into main
4013e1a
@sroy745
Merge branch 'vllm-project:main' into main
2dbdd78
@sroy745
Merge branch 'vllm-project:main' into main
b3575e9
@sroy745
Merge branch 'vllm-project:main' into main
94b0d43
@sroy745
Merge branch 'vllm-project:main' into main
fa8fedf
@sroy745
Merge branch 'vllm-project:main' into main
6ed96b4
@sroy745
Merge branch 'vllm-project:main' into main
b71c533
@sroy745
Merge branch 'vllm-project:main' into main
57babef
@sroy745
Merge branch 'vllm-project:main' into main
4b19bac
@sroy745
Merge branch 'vllm-project:main' into main
eb7a1c4
@sroy745
Merge branch 'vllm-project:main' into main
7e2c87e
@sroy745
Merge branch 'vllm-project:main' into main
6212d5f
@sroy745
Disable LogProbs generation for spec decoding
1eddc27
@github-actions GitHub Actions
Copy link

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only trigger fastcheck CI to run, which consists only a small and essential subset of tests to quickly catch errors with the flexibility to run extra individual tests on top (you can do this by unblocking test steps in the Buildkite run).

Full CI run is still required to merge this PR so once the PR is ready to go, please make sure to run it. If you need all test signals in between PR commits, you can trigger full CI as well.

To run full CI, you can do one of these:

  • Comment /ready on the PR
  • Add ready label to the PR
  • Enable auto-merge.

🚀

@sroy745 sroy745 marked this pull request as draft July 16, 2024 20:45
@sroy745 sroy745 changed the title [WIP] [Spec Decode] Disable Log Probs computation for spec decoding for both draft and target models. [Spec Decode] Disable Log Probs for spec decoding for both draft and target models. Jul 19, 2024
@sroy745 sroy745 changed the title [Spec Decode] Disable Log Probs for spec decoding for both draft and target models. [Spec Decode] Disable Log Prob serialization to CPU for spec decoding for both draft and target models. Jul 19, 2024
@sroy745 sroy745 marked this pull request as ready for review July 19, 2024 19:49
@sroy745
Copy link
Collaborator Author

sroy745 commented Jul 19, 2024

@cadedaniel this pr is now ready for review. PTAL.

cadedaniel
cadedaniel approved these changes Jul 19, 2024
View reviewed changes
Copy link
Collaborator

@cadedaniel cadedaniel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

great stuff, thanks! all feedback is code cleanliness/comments

vllm/spec_decode/spec_decode_worker.py Outdated Show resolved Hide resolved
vllm/spec_decode/spec_decode_worker.py Outdated Show resolved Hide resolved
vllm/spec_decode/target_model_runner.py Outdated Show resolved Hide resolved
vllm/spec_decode/target_model_runner.py


class TargetModelRunner(ModelRunner):
"""Specialized model runner for speculative decoding target model.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: add comment explaining why we do this

In speculative decoding, the logprobs selected may not be the same ones as
selected by the target model sampling. This means that the time spent in the
logprob calculation of the target model is time wasted, since we calculate
logprobs after deciding which tokens are accepted. For this reason we disable
logprobs in the target model so scoring is faster.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

vllm/spec_decode/spec_decode_worker.py Outdated Show resolved Hide resolved
vllm/config.py Outdated Show resolved Hide resolved
@sroy745
Copy link
Collaborator Author

sroy745 commented Jul 21, 2024

Thanks for the review. Addressed all comments. PTAL.

@cadedaniel
Copy link
Collaborator

Enabling auto merge

@cadedaniel cadedaniel merged commit 14f91fe into vllm-project:main Jul 21, 2024
72 checks passed
xjpang pushed a commit to xjpang/vllm that referenced this pull request Jul 24, 2024
@sroy745
[Spec Decode] Disable Log Prob serialization to CPU for spec decoding…
47a7201 
… for both draft and target models. (vllm-project#6485)
gnpinkert pushed a commit to gnpinkert/vllm that referenced this pull request Jul 26, 2024
@sroy745 @gnpinkert
[Spec Decode] Disable Log Prob serialization to CPU for spec decoding…
04eb29f 
… for both draft and target models. (vllm-project#6485)
cduk pushed a commit to cduk/vllm-pascal that referenced this pull request Aug 6, 2024
@sroy745
[Spec Decode] Disable Log Prob serialization to CPU for spec decoding…
a3d1310 
… for both draft and target models. (vllm-project#6485)
@njhill njhill mentioned this pull request Aug 6, 2024
Closed
@tjohnson31415 tjohnson31415 mentioned this pull request Aug 6, 2024
Merged
@tjohnson31415 tjohnson31415 mentioned this pull request Aug 20, 2024
Merged
Alvant pushed a commit to compressa-ai/vllm that referenced this pull request Oct 26, 2024
@sroy745 @Alvant
[Spec Decode] Disable Log Prob serialization to CPU for spec decoding…
3460dee 
… for both draft and target models. (vllm-project#6485)

Signed-off-by: Alvant <[email protected]>
KuntaiDu pushed a commit to KuntaiDu/vllm that referenced this pull request Nov 20, 2024
@sroy745
[Spec Decode] Disable Log Prob serialization to CPU for spec decoding…
626c8cd 
… for both draft and target models. (vllm-project#6485)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@cadedaniel cadedaniel cadedaniel approved these changes

Assignees
No one assigned
Labels
ready ONLY add when PR is ready to merge/full CI is needed
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@sroy745 @cadedaniel