[patrickvonplaten]

What does this PR do?

This PR enables speculative decoding when batch size is > 1.

This PR has been tested to work correctly for Whisper. You can try it as follows:

from transformers import pipeline, AutoModelForCausalLM, AutoModelForSpeechSeq2Seq, AutoProcessor
import torch
from datasets import load_dataset

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

assistant_model_id = "distil-whisper/distil-large-v2"

assistant_model = AutoModelForCausalLM.from_pretrained(
    assistant_model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
assistant_model.to(device)

model_id = "openai/whisper-large-v2"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    generate_kwargs={"assistant_model": assistant_model},
    torch_dtype=torch_dtype,
    chunk_length_s=15,
    batch_size=4,
    device=device,
)

dataset = load_dataset("distil-whisper/librispeech_long", "default", split="validation")
sample = dataset[0]["audio"]

result = pipe(sample)
print(result["text"])

TODOS

• Test that it works correctly for Whisper
• Test for Llama
• Write nice tests
• Clean code
• Write docs

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[MahmoudAshraf97]

I have a question about a part that I don't understand, given that the results are identical with and without batched spec. decoding

this assumes that all the sequences will have the same number of correct matches, when we choose the accepted tokens number based on the maximum correct matches, we also accept wrong tokens from sequences that had lower number of correct matches than the maximum, the solution would be to choose the minimum which guarantees that all the accepted tokens are correct

Suggested change

	valid_tokens = selected_tokens[:, : n_matches.max() + 1]
	valid_tokens = selected_tokens[:, : n_matches.min() + 1]

that poses a risk that if a single sequence had zero correct matches, this whole spec. decoding cycle will be discarded for all other sequences even if they have correct matches.

is it possible to use padding to accept different number of candidate tokens per batch element as opposed to accepting the minimum/maximum common number of correct suggested tokens across the batch, this should at least keep the speedups of spec decoding that were observed at small batch sizes (if not even faster at larger sizer)

[learning-chip]

is it possible to use padding to accept different number of candidate tokens per batch element as opposed to accepting the minimum/maximum common number of correct suggested tokens across the batch

This is possible if the padding lengths are adjusted at every iteration. For one sample that progresses slower, it should be padded more, to keep the same length as fast-progress samples.

[preminstrel]

Hi, good work! I am curious about how to deal with the wrong kv cache token rolling back for draft model, or you just mask them?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[bruno-hays]

@patrickvonplaten Is there any plan to merge this PR sometimes or is it canceled ?

[asartran]

@Hubert-Bonisseur The PR has closed due to inactivity. I'm not sure of the current status of batch > 1 for speculative decoding, but I'm sure @gante will know :)

[bruno-hays]

It is still not supported :
ValueError: assisted generate is only supported for batch_size = 1

But I see that Patrick has left HF since writing this PR, which probably explains why it stalled.
If there is still interest in this PR, I'm willing to help with finish the work of Patrick. It seems like it is very close to be ready and I'm vert interested in this feature.

[bruno-hays]

Maybe your whisper expert @sanchit-gandhi also knows the status of this feature ?

[sanchit-gandhi]

Hey @Hubert-Bonisseur - handing this one over to @kamilakesbi who will take over the integration!

[bruno-hays]

Hey @kamilakesbi, just checking in to see if there are any updates on this. Thanks!