Faster and Smaller Whisper: A Deep Dive into Quantization and Torch Compilation

Table of Contents

Speed Optimization

torch.compile compiles PyTorch models into optimized Triton kernels and can often result in significant speedups for various PyTorch-based models. One of the prerequisites and critical steps toward achieving this is managing a static KV-cache, as detailed in https://pytorch.org/blog/accelerating-generative-ai-2/. During inference, the KV-cache stores intermediate outputs of the model in the attention modules. However, due to varying prompts and generation lengths, these caches are implemented via dynamic allocation, which creates overhead. Instead, torch.compile requires the KV-cache to be maintained statically by pre-allocating the maximum size.

The first step we took was to implement a static cache for the Whisper model to make it fully compatible with torch.compile. A gist of the implementation is available here. As far as we know, this is the first implementation of a static cache for encoder-decoder models in the Transformers library, and hopefully, it will be useful for other encoder-decoder architectures.

Quantization

We further quantized models to 4 bits using Half-Quadratic Quantization (HQQ). HQQ is a fast and accurate model quantizer that does not need any calibration data. Additionally, it fully supports torch.compile. Though quantization techniques can reduce VRAM requirements, they require optimized fused kernels to run faster. Recent developments have introduced low-bit matmul kernels that offer significant speed-ups such as TorchAO's tiny_gemm, and HQQ can leverage these kernels).

Our previous experience with quantizing LLMs has shown that there is only a very marginal drop in quality (measured as word error rate) when quantizing to 4 bits, while achieving further speedup compared to vanilla PyTorch compiled models. Moreover, the GPU VRAM requirements for the decoder can be reduced by a factor of 3 to 4.

We only quantize the decoder linear layers of Whisper for the following reason: the encoder linear layers involve mainly matrix-matrix computations (GEMM), whereas the decoder layers primarily involve matrix-vector multiplications (GEMV) during the decoding phase. The TorchAO kernel is optimized to speed up GEMV operations with 4-bit quantized weights.

Even if the encoder performance is optimized, its impact on the overall runtime is marginal compared to the decoder. The encoder is only run once, whereas the decoder must be run many times, once per token generated.

Benchmarks

We conducted two sets of benchmarking experiments to validate the speed-up using torch.compile and HQQ. The Open ASR eval dataset is used in the first experiment to measure the speed improvement and effect on the WER, and a test set of audio extracts from YouTube videos is used in the second benchmarking. We find that in real-life use cases, transcripts are generally long, hence the second dataset is included. We used the 'large-v2' model with a batch size of 1 and greedy decoding (no sampling) for the experiments.

The Real Time Factor (RTF) metric, commonly used in open-source benchmarks like Open ASR Leaderboard evaluation, measures the speed of offline ASR systems by comparing total processing time to audio duration. We use 1/RTF in our benchmarking experiments as it is more interpretable once the speed becomes faster than the real-time speed, which is a common offline ASR requirement. We used EnglishTextNormalizer() from Open ASR Eval as the text normalizer to compute WER. Experiments use torch nightly build as of 28 May 2024. The experiments are run on a RTX 4090 GPU.

Short-Form Audio

The first dataset represents short-form audio as the average duration of each sample in the dataset is less than 10 seconds. Four Open ASR eval datasets with audios less than 30 seconds are used for the experiments. Baseline system is the current implementation in transformers, that does not use torch.compile operation.

Real-Word Scenario: Long Audios

In order to check the effect of these decoding methods in real-world scenarios with long audios, we used an internal test dataset with 84 minutes duration and verified ground truth. The test set contains 9 audios ranging from 3 minutes to 13 minutes and covers various audio types.

The speed-up in the actual processing time depends on the audio length. For datasets with larger audio lengths, the speed-up is higher as there are more tokens to decode at a higher pace, and the impact of optimization and quantization is more profound.

For example in AMI corpus, which is a corpus with very short audio clips with average length of 2 seconds, the speed up is around 2x compared to baseline versus the long form dataset where the speed up is 4.6x for the HQQ based approach and 3.6x for pytorch compiled one. This speed up can be achieved on top of the fastest sequential /batched execution pytorch based models, resulting in ultra-fast ASR inference.

Tokens per Second

The significant speed-up and impact of optimization and quantization happens in the decoder layers. The following table shows the tokens per second processed by the decoder.

Whisper in Extreme Low-Bit Configurations

Additionally, we tested Whisper in extremely low-bit configurations (i.e., 3, 2, and 1.58 (ternary) bits) with different quantization group sizes. Quantization parameters (namely scale and zero point) are set and computed for a group of entries in a matrix. A lower group size requires more metadata regarding quantization parameters to be stored. Therefore, finding an optimal balance between a large group size and preserving quality is crucial. However, for lower bit sizes, a smaller group size is required (refer to this blog).

The default group size for 4-bit models is 64, so we report results for lower bit sizes at 32 and 16 group sizes. To our surprise, Whisper's performance only degrades slightly and remains quite useful even in extremely low-bit configurations. For 2-bit and 1.58-bit configurations, a group size of 16 is required to be useful. For larger bit sizes, a group size of 64 exhibits comparable accuracy to the non-quantized model.

Ternary weights (1.58 bit) are quite promising since matrix multiplications can be reformulated as additions, potentially resulting in a 70x improvement in efficiency (refer to this paper). We also conducted experiments with pure 1-bit quantization. However, even with a low group size, the Word Error Rate (WER) was very high, and the output was not meaningful.

Note that the PyTorch acceleration kernels we use are optimized for 4 bits, so the speed decreases when the bit size is less than 4 bits. Below, we report the Word Error Rate (WER) along with speed and quantization bitrate and group sizes for these configurations.

The interactive graph below summarizes the performance with different datasets into a scatter plot. Hover or click on a bubble to display the details.

This warrants two future directions we will be pursuing: the development of more optimized kernels for lower bits and the use of post-quantization training, proposed as one of the methods in our 1-bit blog, to recover from low-bit performance degradation.

Caveats

A few caveats to note

1. The TorchAO kernels required modern GPUs.
2. Benchmarking of ASR systems is not fully mature yet. Therefore, these results need to be further battle-tested. However, our internal benchmark with long-form videos is representative of a production scenario, especially for English.

Citation

@misc{sebastian2024whisper1,
title = {Faster and Smaller Whisper: A Deep Dive into Quantization and Torch Compilation},
url = {https://mobiusml.github.io/whisper-static-cache-blog/},
author = {Jilt Sebastian, Husein Zolkepli, Hicham Badri, and Appu Shaji},
month = {May},
year = {2024}
}