[Whisper] Computing features on GPU in batch mode for whisper feature extractor.

[vaibhavagg303]

What does this PR do?

This pull request introduces a feature enabling the computation of audio features for the Whisper model in batch mode on the GPU. This enhancement substantially reduces latency during feature extraction, thereby enhancing overall performance.
cc: @yashjogi

@sanchit-gandhi @ArthurZucker @hollance Please look into it.

fixes #29901

[ArthurZucker]

Very nice addition!

[ArthurZucker]

Looks good. Let's add a test the uses require_torch_gpu to make sure this runs on GPU!

[ArthurZucker]

LGTM, let's ask for @sanchit-gandhi 's approval as well here!

[sanchit-gandhi]

Thanks for this PR @vaibhavagg303 - performance improvements to Whisper are always most welcome! In general, I'm in favour of the changes towards faster feature extractors. Do you have any numbers on how much faster this CUDA variant is? My only reluctancy in merging this PR would be if the performance gains are negligible, but we add extra complexity due to the potential new device argument. To check this, you can adapt and update the toy benchmark that was proposed as part of the first torch stft addition: #26119 (comment). Otherwise, I've left some small suggestions below.

[sanchit-gandhi]

cc @kamilakesbi for viz

[yashjogi]

#29901
The details regarding the time improvement are mentioned in this issue. It reduces the complete inference time for Whisper-small by around 25%, which is significant improvement considering the overall latency of Whisper models. And we also see a significant time decrease in extracting features from 1.5 seconds ( without batch and CPU ) to 0.02 seconds ( by using batch and GPU ).
@sanchit-gandhi @ArthurZucker

[sanchit-gandhi]

Thanks for the benchmark link @vaibhavagg303, that's most helpful! Once follow-up question to your benchmark: you mention that torch gpu stft "cuts the average time for batches of 8 from 1.5 seconds to 0.25 seconds". How does the compute time change for bsz=1, since this is also a common use-case for computing features in short-form mode (<30s).

[sanchit-gandhi]

As a follow-up PR: we'll need to update the ASR pipeline class to compute batched input features to leverage this speed-up, since currently it uses bsz=1 always for the feature extraction

transformers/src/transformers/pipelines/automatic_speech_recognition.py

Line 66 in 9b0a8ea

processed = feature_extractor(chunk, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt")

[yashjogi]

Thanks for the benchmark link @vaibhavagg303, that's most helpful! Once follow-up question to your benchmark: you mention that torch gpu stft "cuts the average time for batches of 8 from 1.5 seconds to 0.25 seconds". How does the compute time change for bsz=1, since this is also a common use-case for computing features in short-form mode (<30s).

This is the average computation time for 8 audios.

Mode	Time
CPU Batch Size 1	1.5 seconds
CPU Batch Size 8	0.25 seconds
GPU Batch Size 8	0.02 seconds

[sanchit-gandhi]

Thanks @yashjogi. To clarify my question above, what's the speed up of CPU Batch Size 1 vs GPU Batch Size 1?

[vaibhavagg303]

Thanks @yashjogi. To clarify my question above, what's the speed up of CPU Batch Size 1 vs GPU Batch Size 1?

It's around 0.02291 seconds for GPU Batch Size 1 @sanchit-gandhi

[yashjogi]

Also, we found out that this is a huge bottleneck for training Whisper -- this simple change reduced the training time by almost 9 times, depending on the CPU resources of the machine we're training on. This means that it would significantly decrease the time to train distil-whisper as well! https://github.com/huggingface/distil-whisper/blob/main/training/run_distillation.py
@sanchit-gandhi

[sanchit-gandhi]

Thanks for the updates @vaibhavagg303 - just one small suggestion about the docstring, otherwise LGTM!

[vaibhavagg303]

Hi @ArthurZucker, can you please review and approve these changes?

[ArthurZucker]

Thanks for getting through this and bringing fast whisper processing! 🤗

[vaibhavagg303]

Thanks, @sanchit-gandhi @ArthurZucker

[sanchit-gandhi]

Thanks for your contribution @vaibhavagg303!