Generate: speculative decoding
Author: gante
Pull request: yes

# What does this PR do?

Useful context:
In a recent PR (#27750), the candidate generation in assisted generation got abstracted, so we can host new candidate generation techniques (such as #27722).

_______________________________________________________

This PR:
1. ~Reworks assisted candidate generation to call `.generate()`, instead of having its own custom generation loop. For most models this is nothing more than a nice abstraction. However, for models with a custom `generate()` function, this means the assistant model will now make use of it! (🤔 does this mean that DistilWhisper gets better numbers with this refactor?)~ Edit: moved to #28030 
2. Adds speculative decoding ([paper](https://arxiv.org/pdf/2211.17192.pdf), see Algorithm 1). This implied a minor interface change in the candidate generation class, which should be okay since it hasn't been released :)

The following tests were run locally and are passing:
1. `RUN_SLOW=1 py.test tests/models/whisper/ -k speculative`
2. `py.test tests/ -k test_assisted` (which now triggers speculative decoding)

________________________________________________________

TODO:
- [ ] Benchmark speculative decoding

Comments:
@patrickvonplaten tagging you here for a 2nd set of eyes on the speculative decoding method (changes in `utils.py`), which I'm assuming you're familiar with. Feel free to delegate to someone else who is familiar with the method! 🤗 

Thanks for adding this! Can we split this up into two separate PRs: one changing the assisted generation and the other adding speculative decoding?

@amyeroberts pulled the assisted generation changes into this PR: https://github.com/huggingface/transformers/pull/28030

After it is merged, I will rebase this one and ping you again -- this one will become exclusively about speculative decoding 🤗 

@amyeroberts I've rerun the slow tests, and I can confirm they are passing. Ready for a review :)

@patrickvonplaten the two types of sampling are needed :D 

New candidate-based methods are popping up (e.g. https://github.com/huggingface/transformers/pull/27775), and they don't necessarily have logits. As such, speculative decoding, which needs the candidates' logits, can't be applied to those methods. 

> @patrickvonplaten the two types of sampling are needed :D
> 
> New candidate-based methods are popping up (e.g. #27775), and they don't necessarily have logits. As such, speculative decoding, which needs the candidates' logits, can't be applied to those methods.

But shouldn't they just be the "own" method now? I.e. I don't think we should put https://github.com/huggingface/transformers/pull/27775 into the speculative decoding method no? 



@patrickvonplaten #27775 does not introduce changes to assisted generation 🤗  In #28030 I've abstracted the candidate generation part of assisted generation. We now load candidate generators the same way as we load the logits processors:

https://github.com/huggingface/transformers/blob/e6dcf8abd6f65bb4b6dfc1831b20d9ba49ce00e2/src/transformers/generation/utils.py#L899-L919

In assisted generation, we call the candidate generator to get candidate sequences (which may or may not contain associated logits, depending on the method)

https://github.com/huggingface/transformers/blob/e6dcf8abd6f65bb4b6dfc1831b20d9ba49ce00e2/src/transformers/generation/utils.py#L4588

The technique in #27775 can thus be added by adding a new candidate generator in ` _get_candidate_generator`. Other candidate generators may be added the same way, enabling users to experiment with the concept of candidates!

Because needing the logits (for speculative decoding) is a very limiting constraint, I'd rather keep the two sampling paths.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_27979). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@amyeroberts PR comments addressed 🤗 

@patrickvonplaten Unless you don't strongly oppose, I'd like to keep the two sampling paths, for the reasons I've written [here](https://github.com/huggingface/transformers/pull/27979#issuecomment-1860125398) -- I think it will be beneficial in the long run! :) (otherwise, a whole new generation method has to be written for #27775)

@amyeroberts -- @patrickvonplaten and I had a chat about whether to keep the two sampling paths or not. For context, here's what we agreed on:
- It's okay to leave it as is, and perhaps abstract the different ways we accept candidates into a `candidate_checker` block. 
- Be conservative on adding new candidate generators, so we don't end up with unused methods
- [in a follow-up PR] squash other cases where the decoding method is the same except for the token selection, like `greedy_decoding` + `sample`
- [in a follow-up PR] mode each decoding method into its own file. There are several private functions in `generation/utils.py` that are exclusively used with one generation method.

@gante 
According to experiments reported in Leviathan's paper, speculative decoding (SD) has higher speedup with greedy decoding (temp=0). However, in the current implementation, SD works only with do_sample=True. 

@jmamou speculative decoding with `do_sample=False` (or `temp=0`) was already encoded in `assisted_generation`, before this PR -- try calling `model.generate(input_ids, do_sample=False, assistant_model=assistant_model)` :)

@gante 
Since acceptance criteria are different between speculative decoding and assisted generation, I think that it would be great to be able to run both speculative decoding and assisted generation with no sampling.

@gante 
I implemented it. I can submit a PR.

@gante 
In previous implementation of assisted generation (4.33) with heuristical update of `num_assistant_tokens` (or `max_assistant_tokens`), the value of `num_assistant_tokens` was preserved between 2 consecutive generate() calls.

In current implementation (4.38), `num_assistant_tokens` is updated by the `candidate_generator` during the generation but `assistant_model.generation_config.num_assistant_tokens` is not updated at the end of the generation. Therefore, next call to generate will start with the initial value of `assistant_model.generation_config.num_assistant_tokens` (5).

Is it intentional? If that's a bug, I can open a PR to fix it. 

@jmamou 

> Since acceptance criteria are different between speculative decoding and assisted generation, I think that it would be great to be able to run both speculative decoding and assisted generation with no sampling.

Not sure if this is a good idea
1. if we see greedy decoding as applying `temperature=0`, the model probability will be `1` at the most likely token and `0` everywhere else. In turn, this implies that `p_i/q_i` is `>=1` at all positions, and thus all candidate tokens would be accepted 👉 speculative decoding would be the same as simply using the assistant model 
2. If we don't apply `temperature=0`, then it would be sampling -- in other words, it wouldn't be greedy decoding

> In previous implementation of assisted generation (4.33) with heuristical update of num_assistant_tokens (or max_assistant_tokens), the value of num_assistant_tokens was preserved between 2 consecutive generate() calls.
In current implementation (4.38), num_assistant_tokens is updated by the candidate_generator during the generation but assistant_model.generation_config.num_assistant_tokens is not updated at the end of the generation. Therefore, next call to generate will start with the initial value of assistant_model.generation_config.num_assistant_tokens (5).
Is it intentional? If that's a bug, I can open a PR to fix it.

This is a good point! A PR to revert to the previous behaviour (with a test) would be appreciated 🙏