[whisper] compile compatibility with long-form decoding
Author: sanchit-gandhi
Pull request: yes

# What does this PR do?

PR #31166 introduced static k/v cache for Whisper short-form decoding. It was noted in this PR that the current generation logic is **not** compatible with sequential long-form generation, since the batch size is reduced dynamically in two places:
1. In the [outer loop](https://github.com/huggingface/transformers/blob/534cbf8a5d332ad8cd26124c5a0a9201f73ca6d7/src/transformers/models/whisper/generation_whisper.py#L639) over time position, we remove audio samples that have already finished generation
2. In the [inner loop](https://github.com/huggingface/transformers/blob/534cbf8a5d332ad8cd26124c5a0a9201f73ca6d7/src/transformers/models/whisper/generation_whisper.py#L817) over temperatures, we remove inputs that don't need fallback at the next temperature increment

For `torch.compile` compatibility with our current cache design, we require the batch size to be fixed. Otherwise, for every batch size we create a new cache object in `.generate`, which changes the data pr of the k/v cache tensors, causing a re-compile.

As things currently stand, we get re-compiles due both the outer and inner loop dynamically changing the batch size. This PR introduces a simple fix: pad the inputs to the max batch size before calling the model, and remove any padded outputs before post-processing.

The alternative would be to change the `batch_idx_map` logic, such that we always keep the full sequence of input features, but only update the sequence generations for the elements of interest. Having tried this quickly as a PoC, the changes are more involved than those proposed in this PR and quickly clutters the dynamic generation logic, which we're retaining for faster eager mode.

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31772). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

cc @ArthurZucker would also be interested in hearing whether you agree with the proposed design

Thanks for the look-over @ArthurZucker. Note that it will increase **average** RAM usage (since we'll always use the max batch size, instead of dropping to lower batch sizes where possible), but it won't increase **total** RAM usage (since we never exceed the max batch size). This is the classic speed-memory trade-off though, and one that is a prerequisite for compile (fixed shapes).

Curious what you mean by ragged inputs - do you have any more details on what this would look like?

This one is ready for final review @ArthurZucker! Note that the failing tests will be addressed by https://github.com/huggingface/transformers/pull/32178, which we'll merge first to `main` and then rebase this PR.

Regarding the docs, there's already a very comprehensive example of using torch compile that includes expected speed-ups: https://github.com/huggingface/transformers/blob/9451a385261b30e7319a2c93285ab76161e8c003/docs/source/en/model_doc/whisper.md?plain=1#L80

I've added one line to make explicit that this is compatible with short and long-form generation: https://github.com/huggingface/transformers/pull/31772/commits/21002ace80194676a44276f9ad1da7793db5354f

@sanchit-gandhi Is there a way the user can specify the cache size, maximum batch size, etc in prefilling the cache for whisper?
Something like:
```
with torch.no_grad(): 
    cache = StaticCache(model.config, batch_size, cache_size, device, compute_dtype)
past_key_values = EncoderDecoderCache(cache, cache)
```
So that this be used in `model.generate()`?


@Jiltseb you can pass `past_key_values` (an instance of `EncoderDecoderCache`) to `generate`, which should do what you described :)

@gante Thanks. However, it seems `model.generate()` does not support `past_key_values` as an instance of `EncoderDecoderCache`. I am getting execution errors when doing so. 
Could you try running it and let me know?

@Jiltseb give me your example, and I'll have a look :) (it's hard to pinpoint an issue without a specific example)

Sure, it's similar to [whisper_compile](https://github.com/sanchit-gandhi/notebooks/blob/main/whisper_compile.ipynb) notebook from @sanchit-gandhi but for different batch sizes:

Here you go:

```
NUM_TOKENS = 100
NUM_WARMUP = 10
NUM_ITERS = 5
ATTN_IMPLEMENTATION = "sdpa"
MODEL_ID = "openai/whisper-medium.en"
BATCH_SIZES = [2,4]
max_batch_size = 4
torch_device = "cuda:0"
device = torch_device 
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
compute_dtype = torch_dtype
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID, attn_implementation=ATTN_IMPLEMENTATION)
model.to(torch_device, dtype=torch_dtype)

is_multilingual = getattr(model.generation_config, "is_multilingual", False)
language = "en" if is_multilingual else None

dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation", trust_remote_code=True)
sample = dataset[30]["audio"]

inputs = processor(audio_sample, sampling_rate=16000, return_tensors="pt").to(torch_device)

with torch.no_grad(): 
    cache = StaticCache(model.config, max_batch_size, 128, device, compute_dtype)
past_key_values = EncoderDecoderCache(cache, cache)

with torch.no_grad():
    for _ in range(NUM_ITERS):
        torch.cuda.synchronize()
        input_features = inputs.input_features.to(torch_dtype)
        input_features_batch = input_features.repeat(batch_size, 1, 1)
        start = datetime.datetime.now()
        out = model.generate(input_features_batch, min_new_tokens=NUM_TOKENS, max_new_tokens=NUM_TOKENS, language=language, begin_suppress_tokens=None, suppress_tokens=None)
        torch.cuda.synchronize() 
        runtime = datetime.datetime.now() - start 
        print(f"Inference took {runtime} seconds")
        tok_per_s = (out.shape[1] * 1) / runtime.total_seconds()
        print(f"Dynamic bsz {1} - {tok_per_s} tok/s")

    model.config.use_cache=True
    model.forward = torch.compile(model.forward, mode="reduce-overhead",  dynamic=True, fullgraph=True)#(NUM_WARMUP + NUM_ITERS) * NUM_TOKENS
    for batch_size in BATCH_SIZES:
        input_features = inputs.input_features.to(torch_dtype)
        input_features_batch = input_features.repeat(batch_size, 1, 1)

        for _ in range(NUM_WARMUP):
            start = datetime.datetime.now()
            torch.cuda.synchronize()
            out = model.generate(input_features_batch, past_key_values=past_key_values, min_new_tokens=NUM_TOKENS, max_new_tokens=NUM_TOKENS, language=language, begin_suppress_tokens=None, suppress_tokens=None)
            torch.cuda.synchronize()
            print(f"Warmup Inference took {datetime.datetime.now() - start} seconds for batch size {batch_size}")

    #Testing for speed
    for batch_size in BATCH_SIZES:
        for _ in range(NUM_ITERS):
            input_features_batch = input_features.repeat(batch_size, 1, 1)
            torch.cuda.synchronize()
            start = datetime.datetime.now()
            out = model.generate(input_features_batch, past_key_values=past_key_values, min_new_tokens=NUM_TOKENS, max_new_tokens=NUM_TOKENS, language=language, begin_suppress_tokens=None, suppress_tokens=None)
            torch.cuda.synchronize() 
            runtime = datetime.datetime.now() - start 
            print(f"Inference took {runtime} seconds for batch size {batch_size}")
            tok_per_s = (out.shape[1] * batch_size) / runtime.total_seconds()
            print(f"Static bsz {batch_size} - {tok_per_s} tok/s")


```

@Jiltseb the example crashes because `EncoderDecoderCache` needs to be initialized with the exact batch size used at inference time, not the maximum batch size it accepts :)

I understand our variable naming is not precise here. I'm going to rename `EncoderDecoderCache`'s `max_batch_size` to `batch_size`.

I thought the `StaticCache` object takes `max_batch_size` argument as per the [doc](https://huggingface.co/docs/transformers/main/en/llm_optims?static-kv=advanced+usage%3A+control+Static+Cache#static-kv-cache-and-torchcompile). Isn't it?

I meant `max_batch_size` for the `StaticCache`, updated the script :)

@Jiltseb you're going to the correct source, but the source has incorrect information -- that's precisely what I will be updating asap :) It has to be the exact batch size.

For further information: while it is technically feasible to accept smaller batch sizes, we would be slicing the cache at each layer. Slicing is an expensive operation, and the point of `StaticCache` is to be fast 🤗 

@gante It still crashes even if I provide a fixed `batch_size`. Have you checked if it works?

By the way, shouldn't it be possible to allocate variable batch size if you do the following:

- predefine a max `batch_size` and preallocate the cache
- when the `batch_size` changes, only the first N elements in the cache are used. 
- warm-up with for range(1, batch_size+1) separately when compiling. 

This can make sure that it runs faster for all batch sizes up to batch_size, especially useful if the actual batch_size is not reached in the long-form audio inference.


@Jiltseb the code is working as intended. A suggestion to avoid waiting on our (sometimes slow) replies: check the documentation of the functions and run without compilation for informative exceptions :)

In particular, you have to make sure the cache for the cross-attention is properly set, as documented in the `EncoderDecoderCache` class. You'll also need to manually reset the cache if you're handling it outside `generate`. See the example below

```py
from transformers import WhisperForConditionalGeneration, AutoProcessor, StaticCache, EncoderDecoderCache
from datasets import load_dataset
import torch
import datetime

NUM_TOKENS = 100
NUM_WARMUP = 10
NUM_ITERS = 5
ATTN_IMPLEMENTATION = "sdpa"
MODEL_ID = "openai/whisper-medium.en"
BATCH_SIZES = [2,4]

torch_device = "cuda:0"
device = torch_device
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
compute_dtype = torch_dtype
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID, attn_implementation=ATTN_IMPLEMENTATION)
model.to(torch_device, dtype=torch_dtype)

is_multilingual = getattr(model.generation_config, "is_multilingual", False)
language = "en" if is_multilingual else None

dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation", trust_remote_code=True)
sample = dataset[30]["audio"]["array"]

inputs = processor(sample, sampling_rate=16000, return_tensors="pt").to(torch_device)


with torch.no_grad():
    model.config.use_cache=True
    model.forward = torch.compile(model.forward, mode="reduce-overhead",  dynamic=True, fullgraph=True)#(NUM_WARMUP + NUM_ITERS) * NUM_TOKENS
    for batch_size in BATCH_SIZES:

        with torch.no_grad():
            cache_cross_attn = StaticCache(model.config, batch_size, model.config.max_source_positions, device, compute_dtype)
            cache_self_attn = StaticCache(model.config, batch_size, 128, device, compute_dtype)
        past_key_values = EncoderDecoderCache(cache_self_attn, cache_cross_attn)

        input_features = inputs.input_features.to(torch_dtype)
        input_features_batch = input_features.repeat(batch_size, 1, 1)

        for _ in range(NUM_WARMUP):
            start = datetime.datetime.now()
            torch.cuda.synchronize()
            out = model.generate(input_features_batch, past_key_values=past_key_values, min_new_tokens=NUM_TOKENS, max_new_tokens=NUM_TOKENS, language=language, begin_suppress_tokens=None, suppress_tokens=None)
            torch.cuda.synchronize()
            print(f"Warmup Inference took {datetime.datetime.now() - start} seconds for batch size {batch_size}")

            past_key_values.reset()
```