Серия интервью с AI #4: объясните, что такое KV-кэширование

Вопрос:
Вы внедряете большую языковую модель в производственную среду. Генерация первых нескольких токенов происходит быстро, но по мере увеличения последовательности генерация каждого дополнительного токена занимает всё больше времени, даже если архитектура модели и аппаратные средства остаются прежними.

Если вычисления не являются основным узким местом, то какая неэффективность вызывает такое замедление, и как можно перепроектировать процесс логического вывода, чтобы генерация токенов стала значительно быстрее?

Что такое KV-кэширование и как оно ускоряет генерацию токенов?

KV-кэширование — это метод оптимизации, используемый при генерации текста в больших языковых моделях для избежания повторных вычислений. В авторегрессивной генерации модель создаёт текст по одному токену за раз, и на каждом шаге обычно заново вычисляет внимание ко всем предыдущим токенам. Однако ключи (K) и значения (V), вычисленные для более ранних токенов, никогда не меняются.

При использовании KV-кэширования модель сохраняет эти ключи и значения при первом их вычислении. При генерации следующего токена она повторно использует кэшированные K и V вместо того, чтобы вычислять их заново, и только вычисляет запрос (Q), ключ и значение для нового токена. Затем внимание рассчитывается с использованием кэшированной информации плюс нового токена.

Такое повторное использование прошлых вычислений значительно сокращает избыточную работу, делая вывод быстрее и эффективнее — особенно для длинных последовательностей — за счёт дополнительной памяти для хранения кэша.

Оценка влияния KV-кэширования на скорость вывода

В этом коде мы оцениваем влияние KV-кэширования во время авторегрессивной генерации текста. Мы запускаем один и тот же запрос через модель несколько раз: один раз с включённым KV-кэшированием и один раз без него, и измеряем среднее время генерации. Поддерживая модель, запрос и длину генерации постоянными, этот эксперимент показывает, как повторное использование кэшированных ключей и значений значительно сокращает повторные вычисления внимания и ускоряет вывод.

«`python
import numpy as np
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

device = «cuda» if torch.cuda.is_available() else «cpu»

model_name = «gpt2-medium»
tokenizer = AutoTokenizer.frompretrained(modelname)
model = AutoModelForCausalLM.frompretrained(modelname).to(device)

prompt = «Объясните, что такое KV-кэширование в трансформерах.»

inputs = tokenizer(prompt, return_tensors=»pt»).to(device)

for use_cache in (True, False):
times = []
for _ in range(5):
start = time.time()
model.generate(
inputs,
usecache=usecache,
maxnewtokens=1000
)
times.append(time.time() — start)

print(
f»{‘с’ if use_cache else ‘без’} KV-кэширования: «
f»{round(np.mean(times), 3)} ± {round(np.std(times), 3)} секунд»
)
«`

Результаты наглядно демонстрируют влияние KV-кэширования на скорость вывода. При включённом KV-кэшировании генерация 1000 токенов занимает около 21,7 секунды, тогда как отключение KV-кэширования увеличивает время генерации до более чем 107 секунд — почти в 5 раз медленнее.

Эта резкая разница возникает потому, что без KV-кэширования модель пересчитывает внимание ко всем ранее сгенерированным токенам на каждом шаге, что приводит к квадратичному росту вычислений. С KV-кэшированием прошлые ключи и значения повторно используются, устраняя избыточную работу и сохраняя время генерации практически линейным по мере роста последовательности. Этот эксперимент подчёркивает, почему KV-кэширование необходимо для эффективного использования больших языковых моделей в реальных условиях.

1. В чём заключается основная проблема при генерации текста в больших языковых моделях без использования KV-кэширования?

Основная проблема заключается в том, что при генерации текста по одному токену за раз модель заново вычисляет внимание ко всем предыдущим токенам на каждом шаге. Это приводит к квадратичному росту вычислений и замедлению процесса генерации, особенно для длинных последовательностей.

2. Как KV-кэширование помогает ускорить генерацию токенов в больших языковых моделях?

KV-кэширование ускоряет генерацию токенов за счёт повторного использования ключей (K) и значений (V), вычисленных для более ранних токенов. Модель сохраняет эти значения при первом вычислении и использует их при генерации следующего токена, вместо того чтобы вычислять их заново. Это сокращает избыточную работу и делает вывод быстрее и эффективнее.

3. Какие результаты демонстрирует эксперимент с использованием KV-кэширования и без него?

Эксперимент показывает, что при включённом KV-кэшировании генерация 1000 токенов занимает около 21,7 секунды, тогда как отключение KV-кэширования увеличивает время генерации до более чем 107 секунд — почти в 5 раз медленнее. Это подчёркивает важность KV-кэширования для эффективного использования больших языковых моделей в реальных условиях.

4. Какие выводы можно сделать из эксперимента с KV-кэшированием?

Из эксперимента можно сделать вывод, что KV-кэширование является необходимым для эффективного использования больших языковых моделей в реальных условиях. Оно позволяет значительно сократить время генерации текста, особенно для длинных последовательностей, за счёт устранения избыточной работы.

5. Какие преимущества даёт использование KV-кэширования при работе с большими языковыми моделями?

Преимущества использования KV-кэширования включают:
* значительное сокращение времени генерации текста;
* устранение избыточной работы;
* повышение эффективности использования больших языковых моделей в реальных условиях;
* сохранение времени генерации практически линейным по мере роста последовательности.

Источник