Dynamic number of speculative tokens in order to accelerate speculative decoding

[jmamou]

What does this PR do?

This PR adds support to dynamic number of speculative tokens in order to accelerate speculative decoding.
It is an unsupervised version of the dynamic speculation lookahead from Dynamic Speculation Lookahead Accelerates Speculative Decoding of Large Language Models.
 
We add the argument assistant_confidence_threshold to the generation configuration of the model. It is a confidence threshold for the assistant model. If the assistant model's confidence in its prediction for the current token is lower than this threshold, the assistant model stops the current token generation iteration, even if the number of speculative tokens (defined by num_assistant_tokens) is not yet reached.

Before submitting

• [X ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• [ V] Did you read the contributor guideline,
  Pull Request section?
• [ X] Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• [ V] Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• [ V] Did you write any new necessary tests?

Who can review?

@gante @amyeroberts

[gante]

Thank you for opening the PR @jmamou 🤗 The idea in the paper is really cool, much better than my scrappy heuristic! 🔥

I've added a few nits in the PR review that should be trivial to solve.

A proposal for a follow-up PR: Different flavors of speculative decoding/assisted generation are (mostly) invisible to the user other than execution speed. According to your paper, DISCO should be faster than the current default. Do you think you can find a single default value for assistant_confidence_threshold such that it beats the current default in most situations? If so, I'd be more that happy to make DISCO the default 🤗

[jmamou]

Thank you for opening the PR @jmamou 🤗 The idea in the paper is really cool, much better than my scrappy heuristic! 🔥

I've added a few nits in the PR review that should be trivial to solve.

A proposal for a follow-up PR: Different flavors of speculative decoding/assisted generation are (mostly) invisible to the user other than execution speed. According to your paper, DISCO should be faster than the current default. Do you think you can find a single default value for assistant_confidence_threshold such that it beats the current default in most situations? If so, I'd be more that happy to make DISCO the default 🤗

Thanks @gante for your feedback.

Concerning the follow-up PR, I have run experiments with vicuna-13b/vicuna-68m as target/assistant on the datasets Alpaca and CNN-DM with do_sample True and False. The values of assistant_confidence_threshold=0.4 and num_assistant_tokens=20 are consistently almost optimal.

[gante]

Thank you for iterating! To make our CI happy, run make fixup and push the changes :)

[jmamou]

Thank you for iterating! To make our CI happy, run make fixup and push the changes :)

sure!
I am done, except for the test "examples_tensorflow" that does not seem to be related to my code.

[jmamou]

Thank you for iterating! To make our CI happy, run make fixup and push the changes :)

sure!
I am done, except for the test "examples_tensorflow" that does not seem to be related to my code.

[LysandreJik]

Thank you for the contribution! The failing test is unrelated to this PR and should now be fixed on main

[gante]

The values of assistant_confidence_threshold=0.4 and num_assistant_tokens=20 are consistently almost optimal.

@jmamou feel free to open a PR to update the defaults, I'd be glad to accept it! 🤗