Fix `_speculative_sampling` implementation
Author: ofirzaf
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Current implementation of `_speculative_sampling` accepts the draft model tokens all the time due to faulty test of the number of matches (`n_matches`). After fixing this issue I found and fixed several more issues in the implementation to reproduce the exact algorithm presented in the paper.



## Before submitting
- [ ] 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#start-contributing-pull-requests),
      Pull Request section?
- [ ] 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.
- [ ] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@gante
@echarlaix 
<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
I was thinking about changing `max_new_tokens` inside the `get_candidates` method so we won't generate more tokens than `max_length - 1`:
```python
max_new_tokens = min(int(self.num_assistant_tokens), self.generation_config.max_length - new_cur_len - 1)
```
This will help avoid checking in several places if `n_matches` is too big, it can also cause a problem in `_speculative_sampling`, and prevent unnecessary compute. 

What do you think?


@ofirzaf good suggestion, the candidate model should indeed be set up such that the resulting candidate sequence has at most a length of `max length - 1`, as the main model will always generate one additional token 👍 

> Thanks for fixing this!
> 
> +1 on Joao's comment reworking to make more understandable and adding epsilon.
> 
> Could you also add at least one test which fails on main but fixes with this fix? It's fine if it's an integration test checks generated outputs

I've been trying to verify the correctness of the Speculative Decoding (SD) implementation. To do so, I'm planning to add a test which verifies that the token distributions we get using SD are indeed similar (or very close) to the token distributions when using standard sample based decoding. If I understand correctly, this is what the SD paper guarantees. 

So far, I found that when comparing the token scores of the above methods on a sample input, they are very similar but not identical, perhaps only due to hardware numerics errors.
Since tokens with very low probability do not affect the sampling outcome, I think the test should verify that the top `k` token scores of both methods are close enough to each other (for some chose `k` like 5-10). Do you think there's a better way?

When I benchmark Assisted Generation vs SD, I get very similar latencies, even though SD should theoretically be faster, due to a more relaxed acceptance criteria. In consequence, I think it's important to verify the correctness of the SD implementation e.g. by comparing outputs like above. 

@gante @ofirzaf 
WDYT?

My results on A100:
```bash
Target: bigcode/starcoder
Assistant: bigcode/tiny_starcoder_py

Data: openai_humaneval, 20 random examples (args.seed=42)

Setup: 
args.seed=42
max_new_tokens=128
output_scores=True

======================================================================================================================================================
Method               Token Latency        Acceptance Rate      Gen. Args
======================================================================================================================================================

sample               45.19ms              ------               {'do_sample': True, 'temperature': 0.2, 'assistant_model': False}
sd                   25.53ms              67.81%               {'do_sample': True, 'temperature': 0.2, 'assistant_model': True}
ag                   24.04ms              70.78%               {'assistant_model': True}
======================================================================================================================================================
```

@amyeroberts I am not sure what kind of test do you expect to see. I wrote the following test which tests `_speculative_sampling` directly with dummy input data. Is this the kind of testing you are looking for?
```python
import unittest    
from transformers.generation.utils import _speculative_sampling
import torch

class Test(unittest.TestCase):
    def test_speculative_sampling(self):
        # assume vocab size 10, input length 5 + 3 generated candidates
        candidate_input_ids = torch.tensor([[8, 0, 3, 9, 8, 1, 4, 5]])  # input tokens
        candidate_logits = torch.tensor([[
            [-10., 10., -10., -10., -10., -10., -10., -10., -10., -10.],  # generated 1
            [-10., -10., -10., -10., 10., -10., -10., -10., -10., -10.],  # generated 4
            [-10., -10., -10., -10., -10., 10., -10., -10., -10., -10.],  # generated 5
            ]])
        candidate_length = 3
        inf = float('inf')
        new_logits = torch.tensor([[
            [-10., 10., -10., -10., -10., -10., -10., -10., -10., -10.],  # accepts 1
            [-10., -10., -10., -10., 10., -10., -10., -10., -10., -10.],  # accepts 4
            [-inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, 10., -inf],  # rejects 5, accepts 8
            [-10., -10., -10., -10., -10., -10., -10., -10., -10., -10.],  # N/A
        ]])
        last_assistant_token_is_eos=False
        max_matches = 5
        validated_tokens, n_matches = _speculative_sampling(
            candidate_input_ids,
            candidate_logits,
            candidate_length,
            new_logits,
            last_assistant_token_is_eos,
            max_matches,
        )
        self.assertTrue(n_matches.item() == 2)
        self.assertTrue(validated_tokens.tolist()[0] == [1, 4, 8])
```

@danielkorat we want to test that the probability to sample token $t$ from the target model is the same as the probability to sample $t$ using speculative. Maybe the following test will work, given a prompt and a target model, calculate the distribution for the next token by simply running the model on the prompt and applying `softmax` on the resulted logits. Then, we can generate $n$ guesses for the next token using speculative. If speculative is implemented correctly and $n$ is large enough you should get a distribution that would match the calculated distribution from before. 

Does that make sense? WDYT?

@ofirzaf Yes we can either do that or compare the logits using `np.allclose()`.
I wonder though if the guarantee holds for the next token only and not for all subsequent tokens.
If it holds only for the next token, we can have the test feed the SD algo with standard sampled (STD) input_ids one at a time (max_new_tokens=1) and compare outputs (one at a time) with STD in one of the suggested methods above.
________________________________
From: Ofir Zafrir ***@***.***>
Sent: Thursday, January 18, 2024 3:16:46 AM
To: huggingface/transformers ***@***.***>
Cc: Korat, Daniel ***@***.***>; Mention ***@***.***>
Subject: Re: [huggingface/transformers] Fix `_speculative_sampling` implementation (PR #28508)


@danielkorat<https://github.com/danielkorat> we want to test that the probability to sample token $t$ from the target model is the same as the probability to sample $t$ using speculative. Maybe the following test will work, given a prompt and a target model, calculate the distribution for the next token by simply running the model on the prompt and applying softmax on the resulted logits. Then, we can generate $n$ guesses for the next token using speculative. If speculative is implemented correctly and $n$ is large enough you should get a distribution that would match the calculated distribution from before.

Does that make sense? WDYT?

—
Reply to this email directly, view it on GitHub<https://github.com/huggingface/transformers/pull/28508#issuecomment-1897597109>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AH26TASLOFOHUY467DYRXPTYPBZX5AVCNFSM6AAAAABB3HU4LOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQOJXGU4TOMJQHE>.
You are receiving this because you were mentioned.Message ID: ***@***.***>


@ofirzaf your test suggestion would be a great addition 👍 

@danielkorat If the claim in the paper is correct, then @ofirzaf suggestion should work regardless of the distance of the target token to the end of the prompt. In other words, we should see the same distribution in the 1st generated token, in the 2nd generated token given the 1st generated token, and so on.

My suggestion would be to generate ~10 tokens for a given prompt ~1M times using super small models (e.g. pythia), to confirm it empirically.

@gante @amyeroberts I have pushed all the changes we discussed. Can you please take a look?

@gante @ofirzaf 
I wrote this code to check our assumptions about the distributions:

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

def top_k(a, k):
    sorted, indices = a.sort(descending=True)
    return sorted[:k], indices[:k]
def get_dist(outputs, num_samples=10**6, norm_factor=1e7):
    probs = outputs.scores[0][0].softmax(dim=-1)
    print("probs:", top_k(probs, k=5))
    next_std_token = torch.multinomial(probs, num_samples=num_samples, replacement=True)
    counts = torch.bincount(next_std_token, minlength=probs.shape[0]).float()
    print("counts:", top_k(counts, k=5))
    norm_counts = counts.float() / norm_factor
    print("softmaxed_counts:", top_k(counts.softmax(dim=-1), k=5)[0])
    print("norm_counts:", top_k(norm_counts, k=5))
    softmaxed_norm_counts = top_k(norm_counts.softmax(dim=-1), k=5)[0]
    print("softmaxed_norm_counts:", softmaxed_norm_counts)
    return softmaxed_norm_counts

checkpoint = "EleutherAI/pythia-1.4b-deduped"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint)
assistant_model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-160m-deduped")
gen_kwargs = dict(**tokenizer("In other words, we", return_tensors="pt"),
                  pad_token_id=tokenizer.eos_token_id,
                  max_new_tokens=1,
                  do_sample=True,
                  temperature=0.1,
                  return_dict_in_generate=True,
                  output_scores=True)
print("Standard sampling:")
set_seed(0)
std_dist = get_dist(model.generate(**gen_kwargs))
print("\nSpeculative decoding:")
set_seed(0)
sd_dist = get_dist(model.generate(**gen_kwargs, assistant_model=assistant_model))
print(std_dist == sd_dist)
```

In short, I get the token probs for some input for both decoding methods, then I sample 1M times from this prob, and get the token counts (I always look at top 5 scores in all the following computations). As you can see below, the `counts` are very close to each other but not equal of course. If we apply `softmax` to the counts, they produce an insignificant result of  `[1., 0., 0., 0., 0.]` (tested other inputs too). So I normalize the counts and then apply `softmax`, and then the output distributions look similar, however, they are still not equal to each other (final output). 

WDYT? Is this the way to go?
If so, I can extend this test to the next 10 generated tokens as well.

Output::
```bash
Standard sampling:
probs: (tensor([5.9927e-01, 3.8301e-01, 1.7703e-02, 8.4421e-06, 1.5724e-06]), tensor([452, 403, 476, 878, 513]))
counts: (tensor([5.9925e+05, 3.8286e+05, 1.7879e+04, 1.7000e+01, 1.0000e+00]), tensor([452, 403, 476, 878, 513]))
softmaxed_counts: tensor([1., 0., 0., 0., 0.])
norm_counts: (tensor([5.9925e-02, 3.8286e-02, 1.7879e-03, 1.7000e-06, 1.0000e-07]), tensor([452, 403, 476, 878, 513]))
softmaxed_norm_counts: tensor([2.1106e-05, 2.0654e-05, 1.9914e-05, 1.9878e-05, 1.9878e-05])

Speculative decoding:
probs: (tensor([5.9925e-01, 3.8304e-01, 1.7706e-02, 8.4404e-06, 1.5720e-06]), tensor([452, 403, 476, 878, 513]))
counts: (tensor([5.9922e+05, 3.8288e+05, 1.7883e+04, 1.7000e+01, 1.0000e+00]), tensor([452, 403, 476, 878, 457]))
softmaxed_counts: tensor([1., 0., 0., 0., 0.])
norm_counts: (tensor([5.9922e-02, 3.8288e-02, 1.7883e-03, 1.7000e-06, 1.0000e-07]), tensor([452, 403, 476, 878, 457]))
softmaxed_norm_counts: tensor([2.1106e-05, 2.0654e-05, 1.9914e-05, 1.9878e-05, 1.9878e-05])
tensor([False, False, False, False, False])
```

@danielkorat there are tiny fluctuations to be expected: doing the forward pass with different shapes (e.g. 1 token at a time vs all at once, 1 row at a time vs batched) will result in slightly different outputs. I've written about it in more detail [here](https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535).

Factoring in this source of numerical differences, the results do look similar! Thank you for running the experiment 🤗 

> @danielkorat there are tiny fluctuations to be expected: doing the forward pass with different shapes (e.g. 1 token at a time vs all at once, 1 row at a time vs batched) will result in slightly different outputs. I've written about it in more detail [here](https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535).
> 
> 
> 
> Factoring in this source of numerical differences, the results do look similar! Thank you for running the experiment 🤗 

@gante My intention is to integrate this code into the repo tests. WDYT?
Should I do this in a separate PR?

@danielkorat This test requires a lot of compute, so it's not fit for our CI 😅