[OLD] New PR: #35029. [[Universal Speculative Decoding `CandidateGenerator`]]
Author: keyboardAnt
Pull request: yes

Please see the new PR: https://github.com/huggingface/transformers/pull/35029

---

~~_This PR is ready for initial review, though some aspects are still a work-in-progress._~~

# ~~What does this PR do?~~

~~This PR introduces the `UniversalSpeculativeDecodingGenerator` class, enabling speculative decoding for assistants with slightly different tokenizers. The key addition is two logits processors (`LogitsProcessor`) that ensure the assistant generates tokens exclusively from the target vocabulary, maintaining alignment and preserving the target distribution without altering the verification method. Theoretically, it is agnostic to the `do_sample` choice. This avoids issues like #32867 and #33534 and sets the stage for advanced universal speculative decoding techniques (that we are currently researching and have not yet been published).~~

---

## ~~Motivation and Context~~

~~This update resolves prior inconsistencies in speculative decoding caused by misaligned vocabularies. Key benefits include:~~
- ~~Ensuring the assistant generates only tokens present in the target vocabulary.~~
- ~~Lossless preservation of the target distribution.~~
- ~~Compatibility with future speculative decoding advancements.~~

~~This PR is a step toward advancements in [Universal Assisted Generation](https://huggingface.co/blog/universal_assisted_generation), in collaboration with @danielkorat, @Orenpereg, @mosheber, @jmamou, @gante, @lewtun, and @MosheWasserb.~~

---

## Related

Issues:
- #32867
- #33534

PRs:
- https://github.com/huggingface/transformers/pull/34553
- https://github.com/huggingface/transformers/pull/34935 - please review and merge

---

## ~~Dependencies~~

~~No additional dependencies.~~

---

## ~~Before Submitting Checklist~~

- [x] ~~Followed the [contributor guidelines]~~(https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request).
- [x] ~~Add functionality tests.~~
- [ ] ~~Verified adherence to target distribution accuracy.~~
- - [x] ~~Please merge https://github.com/huggingface/transformers/pull/34553.~~
- - [ ] ~~Add more tests.~~
- [ ] ~~Measure speedups and add documentation.~~

---

## ~~Who can review?~~

- ~~Speculative decoding and generation: @gante~~
- ~~Tokenizer alignment: @ArthurZucker~~

Comments:
https://arxiv.org/pdf/2404.09492

Found this paper that attempts to align different vocabularies that works across different LLM families and then creates a projection matrix that projects the different LLM outputs to the same embedding domain.

<img width="406" alt="Screenshot 2024-11-20 at 3 27 40 PM" src="https://github.com/user-attachments/assets/70327ffb-9932-4df5-8d85-46919d57f7be">


> https://arxiv.org/pdf/2404.09492
> 
> Found this paper that attempts to align different vocabularies that works across different LLM families and then creates a projection matrix that projects the different LLM outputs to the same embedding domain.
> 
> <img alt="Screenshot 2024-11-20 at 3 27 40 PM" width="406" src="https://private-user-images.githubusercontent.com/41287729/388296986-70327ffb-9932-4df5-8d85-46919d57f7be.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MzIxNTMzMjgsIm5iZiI6MTczMjE1MzAyOCwicGF0aCI6Ii80MTI4NzcyOS8zODgyOTY5ODYtNzAzMjdmZmItOTkzMi00ZGY1LThkODUtNDY5MTlkNTdmN2JlLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDExMjElMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQxMTIxVDAxMzcwOFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWFkODY3YzE2NGJjMDY4MmMzMDJjOGQ5YTI3ZDBmYWFmMGFhYzcyNmYxODZhNTQwNTA0YWNmYzRlNjg0NDk5OGUmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.9jcKX8m3xwIGgO6ZsKt6gkxUO3a3XBJnCN7EeRNjv2k">

Thanks for sharing, Gaurav.

My takeaways from the paper:
1. ~~Out of the box, without additional training, we can test the PR on the pair of TigerBot and Chinese Alpaca. Their vocabularies differ from each other and have the maximum overlap. The Chinese Alpaca repo (https://github.com/ymcui/Chinese-LLaMA-Alpaca) has more than 18k stars so speeding it up could be an impactful example. What are your thoughts?~~ The TigerBot and Chinese Alpaca model families only offer models with ≥7B params, which might be too large for a drafter.
2. With additional training we can do some cool stuff. :)



> > https://arxiv.org/pdf/2404.09492
> > Found this paper that attempts to align different vocabularies that works across different LLM families and then creates a projection matrix that projects the different LLM outputs to the same embedding domain.
> > <img alt="Screenshot 2024-11-20 at 3 27 40 PM" width="406" src="https://private-user-images.githubusercontent.com/41287729/388296986-70327ffb-9932-4df5-8d85-46919d57f7be.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MzIxNTMzMjgsIm5iZiI6MTczMjE1MzAyOCwicGF0aCI6Ii80MTI4NzcyOS8zODgyOTY5ODYtNzAzMjdmZmItOTkzMi00ZGY1LThkODUtNDY5MTlkNTdmN2JlLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDExMjElMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQxMTIxVDAxMzcwOFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWFkODY3YzE2NGJjMDY4MmMzMDJjOGQ5YTI3ZDBmYWFmMGFhYzcyNmYxODZhNTQwNTA0YWNmYzRlNjg0NDk5OGUmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.9jcKX8m3xwIGgO6ZsKt6gkxUO3a3XBJnCN7EeRNjv2k">
> 
> Thanks for sharing, Gaurav.
> 
> My takeaways from the paper:
> 
> 1. ~Out of the box, without additional training, we can test the PR on the pair of TigerBot and Chinese Alpaca. Their vocabularies differ from each other and have the maximum overlap. The Chinese Alpaca repo (https://github.com/ymcui/Chinese-LLaMA-Alpaca) has more than 18k stars so speeding it up could be an impactful example. What are your thoughts?~ The TigerBot and Chinese Alpaca model families only offer models with ≥7B params, which might be too large for a drafter.
> 2. With additional training we can do some cool stuff. :)

I have run similar analysis on the model pairs used in [Universal Assisted Generation blog](https://huggingface.co/blog/universal_assisted_generation).

Last columns in the table below represent respectively the overlap percentage of draft vocab w.r.t. to target vocab and the overlap percentage of target vocab w.r.t. to draft vocab.

<google-sheets-html-origin><!--td {border: 1px solid #cccccc;}br {mso-data-placement:same-cell;}-->
target model |target vocab size  |  draft model | draft vocab size | vocab overlap | overlap  D/T % | overlap T/D % |
-- | -- | -- | -- | -- | -- | -- |
codellama/CodeLlama-13b-Instruct-hf | 32016 | bigcode/tiny_starcoder_py | 49152 | 8481 | 26 | 17
mistralai/Mixtral-8x22B-Instruct-v0.1 | 32768 | double7/vicuna-68m | 32000 | 24184 | 74 | 76
google/gemma-2-9b | 256000 | double7/vicuna-68m | 32000 | 30489 | 12 | 95
mistralai/Mixtral-8x22B-Instruct-v0.1 | 32768 | Qwen/Qwen2-0.5B-Instruct | 151646 | 10566 | 32 | 7
meta-llama/Llama-3.1-70B | 128256 | Qwen/Qwen2-0.5B-Instruct | 151646 | 109566 | 85 | 72
microsoft/Phi-3-medium-128k-instruct | 32011 | Qwen/Qwen2-0.5B-Instruct | 151646 | 9588 | 30 | 6


Seeing this error on the latest commit to `universal-speculatie-decoding` - 

```
Traceback (most recent call last):
  File "/disk/universal_assisted_generation/llama_qwen.py", line 18, in <module>
    outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk/anaconda3/envs/uag/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/disk/universal_assisted_generation/transformers/src/transformers/generation/utils.py", line 2191, in generate
    result = self._assisted_decoding(
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk/universal_assisted_generation/transformers/src/transformers/generation/utils.py", line 4298, in _assisted_decoding
    valid_tokens, n_matches = _speculative_sampling(
                              ^^^^^^^^^^^^^^^^^^^^^^
  File "/disk/universal_assisted_generation/transformers/src/transformers/generation/utils.py", line 4450, in _speculative_sampling
    q_i = q[:, torch.arange(candidate_length), new_candidate_input_ids].squeeze(0, 1)
          ~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: index 5 is out of bounds for dimension 0 with size 5
```

This is when using the following models -
```
prompt = "Alice and Bob"
checkpoint = "meta-llama/Llama-3.1-8B-Instruct"
assistant_checkpoint = "Qwen/Qwen2-0.5B-Instruct"
```


**Update:** I added caching and some tests. The current issue is the dimensions of the output logits.

Update: 

I pulled all the recent commits and run the same example as above - using `Qwen2-0.5B` and `Llama-3.1-8B` - and hit the following error

```
Traceback (most recent call last):
  File "/disk/universal_assisted_generation/llama_qwen.py", line 18, in <module>
    outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk/universal_assisted_generation/transformers/src/transformers/generation/utils.py", line 2191, in generate
    result = self._assisted_decoding(
             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk/universal_assisted_generation/transformers/src/transformers/generation/utils.py", line 4247, in _assisted_decoding
    candidate_input_ids, candidate_logits = candidate_generator.get_candidates(input_ids)
                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 782, in get_candidates
    target_logits = self._atm_translator.get_target_logits(candidate_logits)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 620, in get_target_logits
    target_logits_supported_indices: torch.IntTensor = assistant_logits_supported_indices.apply_(
                                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/disk/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 621, in <lambda>
    lambda x: self._assistant_to_target_input_ids[x]
              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^
KeyError: 151646
```

This looks like expected because the vocab size for `qwen2-0.5B` is `128256`
and for assistant_vocab (`llama-3.1-8B`) - `151646`.

Some debugging tells me this a valid error if the `assistant` model is still generating the tokens from its vocab space rather than being confined to the `target` space.

> Update:
> 
> I pulled all the recent commits and run the same example as above - using `Qwen2-0.5B` and `Llama-3.1-8B` - and hit the following error
> 
> ```
> Traceback (most recent call last):
>   File "/disk/universal_assisted_generation/llama_qwen.py", line 18, in <module>
>     outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
>               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
>   File "/disk/universal_assisted_generation/transformers/src/transformers/generation/utils.py", line 2191, in generate
>     result = self._assisted_decoding(
>              ^^^^^^^^^^^^^^^^^^^^^^^^
>   File "/disk/universal_assisted_generation/transformers/src/transformers/generation/utils.py", line 4247, in _assisted_decoding
>     candidate_input_ids, candidate_logits = candidate_generator.get_candidates(input_ids)
>                                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
>   File "/disk/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 782, in get_candidates
>     target_logits = self._atm_translator.get_target_logits(candidate_logits)
>                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
>   File "/disk/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 620, in get_target_logits
>     target_logits_supported_indices: torch.IntTensor = assistant_logits_supported_indices.apply_(
>                                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
>   File "/disk/universal_assisted_generation/transformers/src/transformers/generation/candidate_generator.py", line 621, in <lambda>
>     lambda x: self._assistant_to_target_input_ids[x]
>               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^
> KeyError: 151646
> ```
> 
> This looks like expected because the vocab size for `qwen2-0.5B` is `128256` and for assistant_vocab (`llama-3.1-8B`) - `151646`.
> 
> Some debugging tells me this a valid error if the `assistant` model is still generating the tokens from its vocab space rather than being confined to the `target` space.

Thanks @gauravjain14.
I fixed the code to pass the tests and added more tests.
Could you please retry? Also, please feel free to add your [test](https://github.com/huggingface/transformers/pull/34760#issuecomment-2495423416) to `tests/generation/test_candidate_generator.py`.

@gauravjain14 
Code below works for me

```
  prompt = "Alice and Bob"
  checkpoint = "meta-llama/Llama-3.1-8B-Instruct"
  assistant_checkpoint = "Qwen/Qwen2-0.5B-Instruct"

  assistant_tokenizer = AutoTokenizer.from_pretrained(assistant_checkpoint)
  tokenizer = AutoTokenizer.from_pretrained(checkpoint)
  inputs = tokenizer(prompt, return_tensors="pt")

  model = AutoModelForCausalLM.from_pretrained(checkpoint)
  assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint)

  generate_kwargs = {
      "do_sample": True,
      "assistant_model": assistant_model,
      "assistant_tokenizer": assistant_tokenizer,
      "tokenizer": tokenizer,
  }

  outputs = model.generate(**inputs, **generate_kwargs)
  
  print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
```

To make this easier to review, I've split off a smaller PR (https://github.com/huggingface/transformers/pull/35009) that focuses purely on refactoring the existing code, without introducing the new Universal SD features.

The refactor aims to:
- Reduce technical debt and improve maintainability.
- Prepare the code for the upcoming Universal SD (this PR).
- Make it more approachable for community contributions.

@ArthurZucker, @gante, I’d love your feedback and review when you have a moment. Thanks so much!

This branch has diverged from `main`. To make it easier, I opened a new PR: https://github.com/huggingface/transformers/pull/35029

Thanks @gauravjain14 for spotting it.