Как работает кэширование

Представьте, что вы ведёте беседу с кем-то, и вместо того, чтобы помнить, что он сказал ранее, он каждый раз начинает с чистого листа, когда вы отвечаете. Разве это не было бы медленно и неэффективно?

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

Для предсказания 1000-го токена модели требуется информация из предыдущих 999 токенов. Эта информация представлена в виде матричных умножений по представлениям токенов.

Чтобы предсказать 1001-й токен, вам нужна та же самая информация из предыдущих 999 токенов, а также любая информация из 1000-го токена. Это огромное количество матричных умножений, которые модели приходится вычислять снова и снова для каждого токена!

Кэш ключей и значений (KV-кэш) устраняет эту неэффективность, сохраняя пары kv, полученные из слоев внимания для ранее обработанных токенов. Сохраненные пары kv извлекаются из кэша и переиспользуются для последующих токенов, что исключает необходимость повторного вычисления.

Кэширование следует использовать только для inference. Его включение во время обучения может привести к непредвиденным ошибкам.

Чтобы лучше понять, как и почему работает кэширование, давайте подробнее рассмотрим структуру матриц внимания.

Матрицы внимания

Масштабированное внимание с точечным произведением вычисляется, как показано ниже, для батча размера b, числа заголовков внимания h, длины последовательности на данный момент T, и размерности на каждый заголовок внимания d_head. $$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$$

Матрицы запроса (Q), ключа (K) и значения (V) представляют собой проекции входных эмбеддингов с формой (b, h, T, d_head).

Для каузального внимания маска не позволяет модели обращать внимание на будущие токены. Как только токен обработан, его представление никогда не меняется в отношении будущих токенов, что означает, что $K_{\text{past}}$ и $V_{\text{past}}$ могут быть закэшированы и переиспользованы для вычисления представления последнего токена. $$\text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{cached}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{cached}}, v_{t}])$$

На этапе инференса для вычисления представления достаточно запроса последнего токена $x_t$ , предсказывающего следующий токен $ t+1 $. На каждом шаге новые векторы ключа и значения сохраняются в кэше и добавляются к прошлым ключам и значениям. $$K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$$

Внимание вычисляется независимо в каждом слое модели, а кэширование выполняется для каждого слоя отдельно.

См. таблицу ниже, чтобы сравнить, как кэширование повышает эффективность.

Класс кэша

Базовый интерфейс кэша KV принимает тензоры ключей и значений для текущего токена и возвращает обновлённые тензоры K и V. Это управляется внутренне методом forward модели.

new_K, new_V = cache.update(k_t, v_t, layer_idx)
attn_output = attn_layer_idx_fn(q_t, new_K, new_V)

При использовании класса Cache из Transformers модуль само-внимания выполняет несколько ключевых шагов для интеграции прошлой и текущей информации.

1. Модуль внимания объединяет текущие пары kv с прошлыми парами kv, хранящимися в кэше. Это создаёт веса внимания с формой (new_tokens_length, past_kv_length + new_tokens_length). Текущие и прошлые пары kv фактически объединяются для вычисления оценок внимания, что обеспечивает осведомлённость модели о предыдущем контексте и текущем вводе.

2. Когда метод forward вызывается итеративно, крайне важно, чтобы форма маски внимания совпадала с общей длиной прошлых и текущих пар kv. Маска внимания должна иметь форму (batch_size, past_kv_length + new_tokens_length). Обычно это обрабатывается внутри generate(), но если вы хотите реализовать собственный цикл генерации с Cache, имейте это в виду! Маска внимания должна содержать значения прошлых и текущих токенов.

Реализация хранилища кэша

Кэши структурированы как список слоёв, где каждый слой содержит кэш ключей и кэш значений. Кэши ключей и значений — это тензоры с формой [batch_size, num_heads, seq_len, head_dim].

Слои могут быть разных типов (например, DynamicLayer, StaticLayer, StaticSlidingWindowLayer), что в основном влияет на обработку длины последовательности и обновление кэша.

Самый простой вариант — это DynamicLayer, который расширяется по мере обработки новых токенов. Размер измерения длины последовательности (seq_len) увеличивается с каждым новым токеном:

cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)

Другие типы слоёв, такие как StaticLayer и StaticSlidingWindowLayer, имеют фиксированную длину последовательности, которая задаётся при создании кэша. Это делает их совместимыми с torch.compile. В случае StaticSlidingWindowLayer существующие токены вытесняются из кэша при добавлении нового токена.

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

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from accelerate import Accelerator

device = Accelerator().device

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map=device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

past_key_values = DynamicCache(config=model.config)
messages = [{"role": "user", "content": "Hello, what's your name."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

generated_ids = inputs.input_ids
max_new_tokens = 10

for _ in range(max_new_tokens):
    outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
    # Greedily sample one next token
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    # Prepare inputs for the next generation step by leaving unprocessed tokens, in our case we have only one new token
    # and expanding attn mask for the new token, as explained above
    attention_mask = inputs["attention_mask"]
    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
    inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
"[INST] Hello, what's your name. [/INST]  Hello! My name is LLaMA,"