Fix _speculative_sampling implementation

[ofirzaf]

What does this PR do?

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

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

[gante]

Three glaring mistakes. I don't know what was on my mind that day, honestly. And worse, my speculative decoding quality checks didn't even pick this up 👀

Thank you very much for catching these bugs and, better yet, for correctly fixing them 💛

[amyeroberts]

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

[ofirzaf]

I was thinking about changing max_new_tokens inside the get_candidates method so we won't generate more tokens than max_length - 1:

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?

[gante]

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

[danielkorat]

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:

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]

@ofirzaf Thanks for iterating - looks great!

Once a test is added, happy to approve. I'll let you and @gante discuss the best way to do this, including @danielkorat's suggestion.

The currently failing tests are unrelated to this PR and a known issue - we're currently trying to fid a fix.

[ofirzaf]

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

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

[ofirzaf]

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

[danielkorat]

@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<#28508 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AH26TASLOFOHUY467DYRXPTYPBZX5AVCNFSM6AAAAABB3HU4LOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQOJXGU4TOMJQHE>. You are receiving this because you were mentioned.Message ID: ***@***.***>

[gante]

@ofirzaf your test suggestion would be a great addition 👍

[gante]

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

[ofirzaf]

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

[danielkorat]

@gante @ofirzaf
I wrote this code to check our assumptions about the distributions:

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

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

[gante]

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

Factoring in this source of numerical differences, the results do look similar! Thank you for running the experiment 🤗

[gante]

@ofirzaf Thank you for iterating with us 💛

@amyeroberts I'm happy with the PR, ready for your final check

[danielkorat]

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

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?

[gante]

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

[amyeroberts]

Thanks for all the work adding this, iterating to make it a mean, clean refactor and adding tests ❤️

[amyeroberts]

Nice :D