Generate: using k-v cache is faster but no difference to memory usage
Hello!I’m benchmarking inference performance using Whisper and the.generate()method, switching between using/not using thek-v cache).My understanding is that when using the cache, inference should be faster (since we don’t recompute k-v states and cache them instead), but VRAM usage higher (since we keep the cached tensors in memory).However, I’m finding that when using cache that inference is faster, but VRAM stays the sameHere are my results with/without cache for the tiny and base Whisper checkpoints:Inf time withInf time withoutVRAM withVRAM withouttiny9.012.013811381base11.318.415231523These experiments are run with greedy decoding, batch size of 1 and 73 eval samples on a 16GB V100. I’m computing VRAM by callingnvidia-smiand monitoring how much usage there is on the GPU.Is this as expected? Or should we see lower VRAM without cache?Notebook:codesnippets/benchmark_whisper_cache.ipynb at main · sanchit-gandhi/codesnippets · GitHubCode snippet to reproduce:from datasets import load_dataset
from transformers import WhisperConfig, WhisperForConditionalGeneration, WhisperProcessor

import torch
from torch.utils.data import DataLoader
import numpy as np

import time
from tqdm import tqdm
import subprocess as sp
import os
import sched

checkpoint_id = "openai/whisper-tiny.en"
processor = WhisperProcessor.from_pretrained(checkpoint_id)

model = WhisperForConditionalGeneration.from_pretrained(checkpoint_id)
model.to("cuda")
model.half()

librispeech = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")

def preprocess(batch):    
    batch["input_features"] = processor(batch["audio"]["array"], sampling_rate=16000, return_tensors="pt").input_features[0]
    return batch

dataset_processed = librispeech.map(preprocess, remove_columns=librispeech.column_names)

dataloader = DataLoader(dataset_processed.with_format("torch"), batch_size=1)


def get_gpu_memory():
    """
    Python equivalent of nvidia-smi, copied from https://stackoverflow.com/a/67722676
    and verified as being equivalent ✅
    """
    output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
    
    COMMAND = "nvidia-smi --query-gpu=memory.used --format=csv"
    
    try:
        memory_use_info = output_to_list(sp.check_output(COMMAND.split(),stderr=sp.STDOUT))[1:]
    
    except sp.CalledProcessError as e:
        raise RuntimeError("command '{}' return with error (code {}): {}".format(e.cmd, e.returncode, e.output))
    
    memory_use_values = [int(x.split()[0]) for i, x in enumerate(memory_use_info)]
    return memory_use_values

# benchmark generation with cache

start = time.time()
for batch in tqdm(dataloader):
    predicted_ids = model.generate(batch["input_features"].to("cuda").half(), max_new_tokens=128, use_cache=True)
runtime = time.time() - start

print("Runtime with: ", runtime)
print("VRAM with: ", get_gpu_memory()[0])

# if we don't delete and re-load the model the GPU use is lower the second time round: warm-up effects?
del model
torch.cuda.empty_cache()

# benchmark without cache

model = WhisperForConditionalGeneration.from_pretrained(checkpoint_id)
model.to("cuda")
model.half()

start = time.time()
for batch in tqdm(dataloader):
    predicted_ids = model.generate(batch["input_features"].to("cuda").half(), max_new_tokens=128, use_cache=False)
runtime = time.time() - start

print("Runtime without: ", runtime)
print("VRAM without: ", get_gpu_memory()[0])Print Output:Runtime with:  8.990428924560547
VRAM with:  1381
Runtime without:  11.993675231933594
VRAM without:  1381Thanks!
2023-02-08T11:56:56Z
Nice write-up!I think the decoder sequence length and the hidden states of the model might be too small to see a difference here in VRAM.The reason VRAM should behigherwhen caching the k,v states is because we cache the projected k,v states of every layer. This means that our cache is of size:2 * (hidden_size) * (num_layers) * (decoder_length)For VRAM computation, this memory is more or less always added to the peak memory of the computation graph.For comparison, we don’t have this memory when not caching. The memory we always have when not caching before doing the attention QK^T computation (which is probs the bottleneck) is 2 * (hidden_size) * 1 * (decoder_length) . Those are the q, v states right that are computed during attention.=> I expect that here (num_layers), (hidden_size) and (decoder_length) are too small to make a difference.The easiest thing to check here would be to use a bigger model and generate to much longer (set eos to None and generate to 256 tokens).
2023-02-08T11:58:02Z
Overall this is an interesting finding though as it means that the k,v cache probably doesn’t play a big role in reducing VRAM for ASR and at that model size.
2023-02-08T13:29:29Z
@sanchit-gandhia few extra numbers – modifying your script to run on GPT-J with FP16 on an 3090, withinput_ids.shape[1]=16andmax_new_tokens=256, we get:14071MBof GPU usage withuse_cache=False13233MBof GPU usage withuse_cache=TrueThe difference becomes more visible with large models and large sequence lengths
2023-02-08T14:21:33Z
Thank you very much for the detailed response!That makes sense that the difference in VRAM with/without using cache is not significant for a model with such low dimensionality.Repeating the experiment with the large-v2 checkpoint (hidden_size=1280, num_layers=32) and generating to 256 tokens yields measurable differences in VRAM, albeit still only marginal:VRAM with: 7597
VRAM without: 7515
Diff: 82(all values in MB)As we expect, the effect is amplified at 512 tokens, scaling (almost) linearly withdecoder_length:VRAM with: 7639
VRAM without: 7519
Diff: 120ASR models tend to generate quite short decoder-lengths. For example, the average token length in the LibriSpeech validation corpus is just~20 tokens. Setting the max length accordingly, we get:VRAM with: 7515
VRAM without: 7511
Diff: 4So pretty insignificant! My intuition is that since VRAM difference with/without cache is proportional to decoder-length, k-v cache doesn’t have a big effect on VRAM for ASR models, even for larger checkpoints.