Speculative decoding speeds up a large model by pairing it with a smaller assistant model. The assistant proposes several tokens, while the larger model verifies those guesses in parallel, preserving its corrections and skipping matching work. section 12.2.2
The technique trades additional memory for lower generation latency:
- Load both models at once, so the deployment needs memory for a large model and its smaller assistant. section 12.2.2
- Use models with the same architecture but different sizes, such as Whisper-Large-v3 and Distil-Whisper-Large-v3. section 12.2.2
- Expect variable gains rather than a guarantee; the described production example reaches roughly a 2× speedup. section 12.2.2
Speech-to-text with Whisper
A Transformers generation setup loads Whisper-Large-v3 as the main model and Distil-Whisper-Large-v3 as the assistant, then passes the assistant through the generation arguments while evaluating an English audio dataset. section 12.2.2
Listing 12.2 Speculative decoding with Whisper
from transformers import (
AutoModelForCausalLM,
AutoModelForSpeechSeq2Seq,
AutoProcessor,
)
import torch
from datasets import load_dataset
from time import perf_counter
from tqdm import tqdm
from evaluate import load
device = "cuda:0" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")
attention = "sdpa"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
model_id = "openai/whisper-large-v3"
assistant_model_id = "distil-whisper/distil-large-v3"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id,
low_cpu_mem_usage=False,
use_safetensors=True,
attn_implementation=attention,
torch_dtype=torch_dtype,
).to(device)
processor = AutoProcessor.from_pretrained(model_id)
assistant_model = AutoModelForCausalLM.from_pretrained(
assistant_model_id,
low_cpu_mem_usage=False,
use_safetensors=True,
attn_implementation=attention,
torch_dtype=torch_dtype,
).to(device)
dataset = load_dataset(
"hf-internal-testing/librispeech_asr_dummy",
"clean",
split="validation",
trust_remote_code=True,
)
wer = load("wer")
generate_kwargs_1 = {
"language": "en",
"task": "transcribe",
}
generate_kwargs_2 = {
"language": "en",
"task": "transcribe",
"assistant_model": assistant_model,
}
spec_decoding = False
for i, generate_kwargs in enumerate([generate_kwargs_1, generate_kwargs_2]):
all_time = 0
predictions = []
references = []
for sample in tqdm(dataset):
audio = sample["audio"]
inputs = processor(
audio["array"],
sampling_rate=audio["sampling_rate"],
return_tensors="pt",
)
inputs = inputs.to(device=device, dtype=torch_dtype)
start_time = perf_counter()
output = model.generate(
**inputs,
**generate_kwargs,
)
gen_time = perf_counter() - start_time
all_time += gen_time
predictions.append(
processor.batch_decode(
output, skip_special_tokens=True, normalize=True
)[0]
)
references.append(processor.tokenizer.normalize(sample["text"]))
score = wer.compute(predictions=predictions, references=references)
if i > 0:
spec_decoding = True
print(f"Speculative Decoding: {spec_decoding}")
print(f"Time: {all_time}")
print(f"Word Error Rate: {score}")
Inspect the complete speculative-decoding example, including model loading, device selection, timing, and evaluation setup. section 12.2.2
