Universal Assisted Generation: Assisted generation with any assistant model (by Intel Labs)
Author: danielkorat
Pull request: yes
Labels: name: run-slow

# 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:
```python
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
- [x] 
This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).

- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] 
Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link to it if that's the case.

- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

@gante (all others)
CC @lewtun 

 -->


Comments:
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:

```python
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)
```

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 ?)

> 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

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](https://github.com/vllm-project/vllm/issues/7252).

@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** |



@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 :)

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

@danielkorat ping me when it's ready for a review 🤗 

@gante I have addressed your comments, feel free to review. 😃
Regarding `AssistantConfig`: I implemented it as you suggested, but it causes a huge slowdown in performance.
I suspect that since it holds two tokenizer objects, and lives inside `GenerationConfig`, it causes this overhead. I'm still trying to mitigate this. I'm open to any suggestion / solution.



@gante 
Update:
I managed to find a workaround for the slowdown bug. In `AssistedCandidateGeneratorDifferentTokenizers.get_candidates()`, I remove `assistant_config` from `assistant_generation_kwargs['generation_config']`, before calling `self.assistant_model.generate()`:

```python
        # 2. Forecast next N tokens using the assistant model.
        assistant_generation_kwargs = {
            self.input_ids_key: assistant_input_ids,
            "min_new_tokens": min_new_tokens,
            "max_new_tokens": max_new_tokens,
            "generation_config": self.generation_config,
            "logits_processor": self.logits_processor,
        }

        self.assistant_kwargs.pop("attention_mask", None)

        assistant_config = self.generation_config.__dict__.pop("assistant_config", None)
        assistant_output = self.assistant_model.generate(**assistant_generation_kwargs, **self.assistant_kwargs)
        self.generation_config.assistant_config = assistant_config
```

 This mitigates the issue. Please advise if you think there's a better solution.

@gante 
Another workaround is to pass only one additional parameter, `assistant_tokenizer`, directly to `generate()` using `kwargs`. Passing `tokenizer` to `generate()` through `kwargs` is already implemented:

https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/generation/utils.py#L1803

WDYT?

@gante 
After running some tests, this line is causing the slowdown mentioned above:

https://github.com/huggingface/transformers/blob/f5247aca0151de9b53096772d3ba0516ee1b896f/src/transformers/generation/candidate_generator.py#L158

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`:

```python
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.

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.



Thanks @gante, applied your suggestions 😃 

Ready to merge (failing tests are unrelated to this PR)



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

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

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

Sorry for the delay, merging! 

Great work, congrats 🔥