Universal Assisted Generation: Assisted generation with any assistant model (by Intel Labs)

[danielkorat]

What does this PR do?

Co-Authors: @mosheber, @jmamou, @orenpereg

Assisted Generation (AG; i.e. Speculative Decoding) can accelerate inference speed by using a small assistant model in tandem with the original (target) model. However, it only supports assistant models that use the same tokenizer as the target model. This significantly limits the usability of AG, as only few models have such corresponding assistant models.
This PR alleviates that requirement, allowing AG execution with any assistant model regardless of its tokenizer.
Thus, AG can support any target model (!), allowing acceleration of a wide range of models that was not possible until now.

For example, our initial experiments show the following speedups:

target	assistant	dataset	task	speedup
codellama/CodeLlama-13b-Instruct-hf	bigcode/tiny_starcoder_py	openai/humaneval	code generation	2.06x
microsoft/Phi-3-medium-128k-instruct	Qwen/Qwen2-0.5B-Instruct	tau/scrolls	long-context summarization	1.63x
google/gemma-2-9b	double7/vicuna-68m	cnn_dailymail	summarization	1.90x

Our method works by decoding the speculated tokens and re-encoding them using the target tokenizer, and then performing an alignment procedure, to account for inter-tokenizer discrepancies.

Usage example:

out = model.generate(input_ids, do_sample=False, assistant_model=assistant_model, 
                                  assistant_tokenizer=assistant_tokenizer, tokenizer=target_tokenizer
)

where target_tokenizer and assistant_tokenizer are the respective tokenizers; target_lookbehind and assistant_lookbehind are optional parameters which define the window size to consider for solving inter-tokenizer discrepancies after every generation step of the assistant model.

Experimental setup:
1 x A6000 GPU
100 examples per dataset

@gante

CC @lewtun

Before submitting

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

[danielkorat]

Note:
The usage can be streamlined by creating a new config class e.g. SpeculativeConfig, which initializes the tokenizers and has good default values for the lookbehinds:

sd_config = SpeculativeConfig(model_name_or_path=model_name_or_path, assistant_name_or_path=assistant_name_or_path)

out = model.generate(input_ids, do_sample=False, max_new_tokens=max_new_tokens, sd_config=sd_config)

[gante]

Quick question before I dive into the code: I see @jmamou as a co-author, is this PR related to DISCO (and, therefore, to #33258 ?)

[jmamou]

Quick question before I dive into the code: I see @jmamou as a co-author, is this PR related to DISCO (and, therefore, to #33258 ?)

no ... it is orthogonal

[danielkorat]

Hi @gante
Our team greatly values your SD and AG implementations in HuggingFace! 😃
We developed this feature because we wanted to solve one of the main pain points of Speculative Decoding, which does not support model pairs from different families.
Our feature is highly desirable by the community, since it supports acceleration of models which were not supported until now. As an example, our feature is requested in vLLM: link.

[danielkorat]

@gante

We even got a significant speedup using a 9 billion target model (google/gemma-2-9b):

target	assistant	dataset	task	speedup
codellama/CodeLlama-13b-Instruct-hf	bigcode/tiny_starcoder_py	openai/humaneval	code generation	1.96x
microsoft/Phi-3-medium-128k-instruct	Qwen/Qwen2-0.5B-Instruct	tau/scrolls	long-context summarization	1.42x
google/gemma-2-9b	double7/vicuna-68m	cnn_dailymail	summarization	1.70x

[gante]

@danielkorat @jmamou apologies for the delayed review, bugfixes in the pipeline 👼 I had a quick look at the PR and it looks really really cool, thank you for proposing this technique and opening the PR 💛

My immediate comment goes in the same direction as a comment added above, about having a SpeculativeConfig AssistantConfig to parameterize the added technique, living inside GenerationConfig. I am suggesting naming it after assisted generation simply because of consistency throughout the library. Let's please add it, and avoid adding new arguments to generate 🤗 This has a few immediate advantages:

1. simpler generate signature (advanced users will look into GenerationConfig, basic users will appreciate not having to read more flags)
2. we gain the ability to serialize a GenerationConfig with a AssistantConfig to automatically enable your technique
3. In a future PR, I am thinking of moving all flags related to assisted generation there, to add some organization to the large number of flags :)

Let me know if I can help in any way beyond reviewing. I've placed this PR on my top priorities :)

[gante]

I think I will have to work next on torch.compile + assisted generation, to squeeze the most of these speedups 😈

[gante]

A few comments and suggestions 🤗

There is a question I want you folks to check first, before all others. I've tagged it with ❗ ❗

[gante]

❗ ❗

Sanity check: you have confirmed that the runtime with this logic is faster than simply re-encoding the whole input_ids?

If it is just a tiny bit faster, like 5%, then I'd prefer having the simpler reencoding code path exclusively -- it's much simpler to maintain and document

(I'll likely have a few comments to make this section more beginner-friendly, but I'm holding them until we confirm this part :D)

[danielkorat]

Yes, we have tested both options and this one is significantly faster.
Natuarlly, as the number of new tokens increases it becomes even more significant.

[danielkorat]

In one test we did, it was 15% faster. And that test used cnn_dailymail which does not even have long prompts.

Re-encoding (decode+encode) the whole input_ids (which increase in size at at each step) at every assistant generation step incurs a significant overhead.

[gante]

awesome, thank you for confirming

(going to leave this comment unresolved for future reference)

[danielkorat]

@gante
After running some tests, this line is causing the slowdown mentioned above:

transformers/src/transformers/generation/candidate_generator.py

Line 158 in f5247ac

self.generation_config = copy.deepcopy(generation_config)

The line above is super slow since it deep-copies two Tokenizer objects living inside generation_config.assistant_config.

A proper solution could be implementing something like this for GenerationConfig:

from copy import deepcopy

class Foo:
    def __init__(self, content, linked_to):
        self.content = content
        self.linked_to = linked_to

    def __deepcopy__(self, memo):
        # create a copy with self.linked_to *not copied*, just referenced.
        return Foo(deepcopy(self.content, memo), self.linked_to)

where generation_config.assistant_config won't get deep-copied.

[gante]

Hi @danielkorat 👋 I got it right from the comments above, passing the assistant tokenizer separately (like we pass the assistant model) is an option -- let's go with it :)

The deep copy is a hard requirement, to prevent side-effects -- we want to make adjustments to the object inside generate, but we don't want to modify the original object.

[gante]

Aside from the tokenizer placement in the generation config, LGTM 🤗

[danielkorat]

Thanks @gante, applied your suggestions 😃

Ready to merge (failing tests are unrelated to this PR)

[gante]

Perfect! Thank you for iterating 🤗

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[gante]

(the broken CI should be fixed by #33950 , let's wait for it to be merged to then rebase :) )

[ArthurZucker]

Very nice addition! Needs a tad bit more documentation!

[danielkorat]

hi @gante @ArthurZucker
is there an ETA for this merge?
Thanks

[ArthurZucker]

Sorry for the delay, merging!

Great work, congrats 🔥