Whisper v-3 pipeline requiring a lot of memory when setting return_timestamps="word" 
Author: ciroantoniomami
Pull request: no

### System Info

- `transformers` version: 4.35.2
- Platform: Linux-5.15.109+-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.0.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = pipeline(task="automatic-speech-recognition",
                                    model='openai/whisper-large-v3',
                                    torch_dtype=torch.float16,
                                    device=device)

audio, _ = librosa.load(local_filename, sr=el.SAMPLE_RATE, dtype=np.float32)

out = model(
                            inputs=full_audio,
                            chunk_length_s=30,
                            batch_size=el.BATCH_SIZE,
                            return_timestamps='word',
                            generate_kwargs={"language": f"<|{language}|>", "task": "transcribe"},
                        )

### Expected behavior

When using batch_size == 24 and return_timestamps=='word' the pipeline requires more than 20gb of GPU Ram, while with
return_timestamps==True it requires less than 7gb

Comments:
Hey! That's pretty much expected you are adding extra computation in order to compute the word alignment using cross attention. We don't really seem to have a lot of documentation on this part IMO cc @sanchit-gandhi. Pretty much all the processing that happens in timestamp processing is not part of the doc 😅 

Hi @ArthurZucker! Thanks a lot for your answer and your time. I'm aware that producing word alignment requires extra work, btw with other implementation of whisper, such as https://github.com/SYSTRAN/faster-whisper I didn't experience any increase in memory usage. This difference is due to the fact that hf pipeline allows batching? Perhaps is possible to free up some space after the computation of a batch? What I notice is that the memory usage increase as the number of batch computed progress. Maybe this could be the input for future improvements, thanks a lot for your work guys!

Oh nice then a memory leak maybe. cc @sanchit-gandhi and @ylacombe if you can have a look! 🤗 

Thanks for raising this issue @ciroantoniomami, I can take a look at it! A last question before I can try to reproduce it though:
Which audio file did you use ? I'm particularly interested by its duration  

@ylacombe I'm usually working with wav mono audio file, loaded with librosa, of duration between one hour and two hours long.

After investigation, my main hypothesis is that the memory explodes for two reasons:

-  in part because computing word-level timestamps requires to keep track of the [attention weights](https://github.com/huggingface/transformers/blob/08a6e7a702d06826659eb7f0f6b9f37d33f31829/src/transformers/models/whisper/modeling_whisper.py#L2225).  Word-level timestamps requires some cross-attentions , so that's not something that we can bypass easily.  This is something that I've verified when using `output_attentions` with `return_timestamps=False`. In that case, memory usage is a few GB higher `(batch_size=24)`.
- and mostly, [some computation on those cross-attention weights](https://github.com/huggingface/transformers/blob/08a6e7a702d06826659eb7f0f6b9f37d33f31829/src/transformers/models/whisper/modeling_whisper.py#L2548-L2557) seems to introduce memory leakage.



I'm not sure what would be the best way to deal with that issue tbh, I've tried refactoring [the lines in cause](https://github.com/huggingface/transformers/blob/08a6e7a702d06826659eb7f0f6b9f37d33f31829/src/transformers/models/whisper/modeling_whisper.py#L2548-L2557) but memory still leaks, any ideas @ArthurZucker, @Narsil and @sanchit-gandhi  ?




Would do some profiling checking the data pointers and whether or not using the garbage cleaner explicitly helps etc. 
Otherwise not sure 😅 

Agree that a deeper memory inspection through a profile is the best next step here

@ylacombe How do you judge it's a leak ? The code you're referring to is creating entire new tensors, which ofc will occupy more memory, it's not a leak exactly.

A leak would occur if that memory wouldn't be freed when you discard all the results (including `generate_outputs`)


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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I ran whisper-v3-large on a 81 minute (EDIT: previously said 30, now I checked it) audio with `return_timestamps="word"`. It didn't fit on my RTX 4090, so I ran it on CPU. It allocated about 50 GiB of RAM.

Is there some neat way to track allocations?

One point to support the leak theory: Even with `return_timestamps=True` (i.e. not "word") and running on my GPU, when running repeatedly in a loop, I eventually run out of VRAM.

If `return_timestamps=True` has a memory leak it would be easier to track down. `return_timestamps="word"` is expected to take a lot more RAM, not sure exactly how much. 
Feel free to open a new issue for `return_timestamps=True`, and ping @ylacombe and @sanchit-gandhi  making sure you are evaluating the model (generate always de-acvtivates gradients so if you use generate it should no be a problem) 

Good evening, any chance something came up? I really would like to use return_timestamps="word" @ylacombe @sanchit-gandhi thanks for your support.

I encountered a similar thing where using the `return_timestamps="word"` parameter during inference led to higher memory consumption compared to using return_timestamps="True". 
Specifically, with the large-v3 model, I found that I could only run inference with a **batch size of 4** when using `return_timestamps="word"`, whereas a **batch size= 16** worked fine with `return_timestamps="True"`. 
I'm unsure if there is a memory leak or expected behaviour. 

@ArthurZucker @ylacombe  @sanchit-gandhi if you have any insights on this, I would appreciate hearing from you.

It is expected 🤗 

Hello! Any news on this?

I've also been running into this issue--based on my testing it seems to be a regression introduced in version `4.42`. A simple solution for now could be to pin `transformers<4.42`.

| torch version | transformers version | success (not OOM) |
| --- | --- | --- |
| 2.4.1 | 4.45.1 | no |
| 2.4.1 | 4.44.2 | no |
| 2.4.1 | 4.43.4 | no |
| 2.4.1 | 4.42.4 | no |
| 2.4.1 | 4.41.2 | yes |
| 2.4.1 | 4.40.2 | yes |

Code:

```python
import sys
import time

import torch

from transformers import WhisperFeatureExtractor, WhisperForConditionalGeneration, WhisperTokenizerFast, pipeline

model_id = "openai/whisper-large-v3-turbo"

torch_dtype = torch.float16
device = "cuda:0"

model = WhisperForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch_dtype,
).to(device)

tokenizer = WhisperTokenizerFast.from_pretrained(model_id)
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_id)

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=tokenizer,
    feature_extractor=feature_extractor,
    torch_dtype=torch_dtype,
    device=device,
)

outputs = pipe(
    "/path/to/two-hour-audio.mp3",
    chunk_length_s=30,
    batch_size=24,
    generate_kwargs={
        "task": "transcribe",
        "language": "en",
    },
    return_timestamps="word",  # Enable word-level timestamps
)

print(outputs["text"])
```

Hey @alexcannan, thanks for giving details about your issue, and a reproducer!

Do you think you'd be able to tell us what's your GPU specs?

I'm not necessarily surprised that your GPU OOM using a BS=24 tbh, but I'd be curious to know if you already almost reached your GPU's limits with `transformers<4.42`



I'm running this on a 3090 with 24,576 MiB of VRAM. `Driver Version: 555.58.02` `CUDA Version: 12.5`

Here's my approximate max memory usage with various settings. Peak is approximated by watching `watch -d -n 0.2 nvidia-smi`

| transformers version | batch_size | VRAM Approximate Peak |
| --- | --- | --- |
| `4.41.2` | 24 | 6818 MiB |
| `4.42.4` | 24 | OOM |
| `4.42.4` | 1 | 4994 MiB |
| `4.42.4` | 2 | 7850 MiB |
| `4.42.4` | 4 | 13690 MiB |
| `4.42.4` | 8 | OOM |
| `4.41.2` | 12 | 4422 MiB |


Thanks for the detail! It's quite surprising to see it happening with BS=8. Are you running the latest version of transformers?

cc @eustlb 

Measuring peak memory usage with `nvidia-smi` is unfortunately highly unreliable, do you think you could get inspiration from the code snippet [here](https://github.com/huggingface/transformers/pull/33732#issuecomment-2520959718). It should help you get a more accurate estimation!



I'm using this relatively simple code to transcribe a wav audio file using whisper v3 turbo and transformers library:
```
model = AutoModelForSpeechSeq2Seq.from_pretrained(
   "openai/whisper-large-v3-turbo", 
   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,
    torch_dtype=torch_dtype,
    device=device,
)

transcription_result = pipe(
    input_file,
    return_timestamps="word",
    generate_kwargs={"language": "polish"},
)
```

I need return_timestamps="word" to get words timestamps and later to merge words+timestamps with diarization from pyannote to create a role-labeled transcription. BTW: This is single audio file that is 6min long and has 6MB.

It works well regardless of transformers version, however:

1) with transformers version 4.42.4 the code uses ~4.9GB of vRAM on my GPU and at the peak around 2GB of additional RAM (increase from 8GB to 10GB), plus it runs for around 30 seconds.

2) but starting with transformers==4.43.0 and up to current 4.48.0 the code uses the same (~4.9GB) amount of vRAM on my GPU, **but RAM consumption on my PC hits the max of 32GB and the whole thing takes like ~130 seconds to finish. The PC is practically unresponsive until the process ends**.

So something definitely happened between 4.42.4 and 4.43.0 version and this option: return_timestamps="word" is basically unusable with transformers>4.42.4.

Is there any chance anybody can check what might be a problem? 
This is definitely not normal behaviour as was suggested above.

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@ylacombe this hasn't been fixed, has it?

This has not been fixed. And I see the exact problem is in the function split_by_batch_index:
            if key in ["scores", "encoder_attentions", "encoder_hidden_states", "logits"]:
                return [v[batch_idx].cpu() for v in values]

Previously it was moving only the scores into the CPU, now it is moving the scores plus the huge encoder attentions. 


![Image](https://github.com/user-attachments/assets/24f960b8-4646-4550-9a49-fe783abd6654)

> Thanks for the detail! It's quite surprising to see it happening with BS=8. Are you running the latest version of transformers?
> 
> cc [@eustlb](https://github.com/eustlb)
> 
> Measuring peak memory usage with `nvidia-smi` is unfortunately highly unreliable, do you think you could get inspiration from the code snippet [here](https://github.com/huggingface/transformers/pull/33732#issuecomment-2520959718). It should help you get a more accurate estimation!

Hey @ylacombe, 

Facing the same vRAM OOM memory issue in the code snippet you provided when transcribing multiple audios in parallel by specifying them as a list and passing batch_size parameter.
The same pattern follows in batches of shorter audios (less than 30s) too.

Transformers==4.49.0
GPU Spec:  A6000,  with 48GB vRAM

Code Snippet:
```
from transformers import set_seed
import torch

from transformers import WhisperFeatureExtractor, WhisperForConditionalGeneration, WhisperTokenizerFast, WhisperProcessor, pipeline
from datasets import load_dataset

model_id = "openai/whisper-large-v3-turbo"

torch_dtype = torch.float16
device = "cuda:7"

model = WhisperForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch_dtype,
    attn_implementation="eager"
).to(device)

processor = WhisperProcessor.from_pretrained(model_id)
tokenizer = WhisperTokenizerFast.from_pretrained(model_id)
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_id)

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

NB_LOOPS = 1

set_seed(0)

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=tokenizer,
    feature_extractor=feature_extractor,
    torch_dtype=torch_dtype,
    device=device,
)

def measure_latency_and_memory_use(model, nb_loops=NB_LOOPS, chunk_length_s=0.0, batch_size=None):
    # define Events that measure start and end of the generate pass
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    # reset cuda memory stats and empty cache
    torch.cuda.reset_peak_memory_stats(device)
    torch.cuda.empty_cache()
    torch.cuda.synchronize()

    # get the start time
    start_event.record()

    # actually generate
    for _ in range(nb_loops):
        # set seed for reproducibility
        set_seed(0)

        outputs = pipe(
            [sample["array"] for i in range(16)],
            chunk_length_s=chunk_length_s,
            batch_size=batch_size,
            generate_kwargs={
                "task": "transcribe",
                "language": "en",
            },
            return_timestamps="word",  # Enable word-level timestamps
        )
    # get the end time
    end_event.record()
    torch.cuda.synchronize()

    # measure memory footprint and elapsed time
    max_memory = torch.cuda.max_memory_allocated(device)
    elapsed_time = start_event.elapsed_time(end_event) * 1.0e-3

    print("Execution time:", elapsed_time / nb_loops, "seconds")
    print("Max memory footprint", max_memory * 1e-9, " GB")

    return outputs


print("--------------------BS = 4 -------------")

print("USING PÏPELINE WHISPER LONGFORM")

outputs = measure_latency_and_memory_use(model, chunk_length_s=0,batch_size=4)
print("---")


print("--------------------BS = 8 -------------")

print("USING PÏPELINE WHISPER LONGFORM")

outputs = measure_latency_and_memory_use(model, chunk_length_s=0,batch_size=8)
print("---")

print("--------------------BS = 16 -------------")

print("USING PÏPELINE WHISPER LONGFORM")

outputs = measure_latency_and_memory_use(model, chunk_length_s=0,batch_size=16)
print("---")
```
Result:
```
--------------------BS = 4 -------------
USING PÏPELINE WHISPER LONGFORM
Execution time: 131.892625 seconds
Max memory footprint 13.884659200000002  GB
---
--------------------BS = 8 -------------
USING PÏPELINE WHISPER LONGFORM
Execution time: 127.338203125 seconds
Max memory footprint 26.098100224000003  GB
---
--------------------BS = 16 -------------
USING PÏPELINE WHISPER LONGFORM
torch.cuda.OutOfMemoryError: CUDA out of memory. 
```

Does not seem to have been fixed, reopening...

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Closing because this is expected