Como o cache funciona

Imagine que você está tendo uma conversa com alguém e, em vez de lembrar do que a pessoa disse anteriormente, ela precisa recomeçar do zero a cada vez que você responde. Isso seria lento e ineficiente, certo?

Você pode estender essa analogia para modelos transformer. A geração em modelos autoregressivos pode ser lenta porque faz uma previsão um token de cada vez. Cada nova previsão depende de todo o contexto anterior.

Para prever o 1000º token, o modelo requer informações dos 999 tokens anteriores. Essas informações são representadas como multiplicações de matrizes entre as representações dos tokens.

Para prever o 1001º token, você precisa das mesmas informações dos 999 tokens anteriores, além de qualquer informação do 1000º token. Isso representa muitas multiplicações de matriz que o modelo precisa calcular repetidamente para cada token!

Um cache de chave-valor (KV) elimina essa ineficiência armazenando pares de kv derivados das camadas de atenção dos tokens processados anteriormente. Os pares de kv armazenados são recuperados do cache e reutilizados para os tokens subsequentes, evitando a necessidade de recálculo.

O uso de cache deve ser restrito à inferência. Ele pode causar erros inesperados se estiver habilitado durante o treinamento.

Para entender melhor como e por que o cache funciona, vamos examinar de perto a estrutura das matrizes de atenção.

Matrizes de atenção

O atento de produto escalar com escala é calculado conforme mostrado abaixo para um lote de tamanho b, número de cabeças de atenção h, comprimento da sequência até o momento T, e dimensão por cabeça de atenção d_head. $$\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V$$

As matrizes de consulta (Q), chave (K) e valor (V) são projeções dos embeddings de entrada, com forma (b, h, T, d_head).

Para a atenção causal, a máscara impede que o modelo preste atenção a tokens futuros. Uma vez que um token é processado, sua representação nunca muda em relação aos tokens futuros, o que significa que $K_{\text{past}}$ e $V_{\text{past}}$ podem ser armazenados em cache e reutilizados para calcular a representação do último token. $$\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}])$$

No momento da inferência, você só precisa da query do último token para calcular a representação $x_t$ que prevê o próximo token $ t+1 $. A cada etapa, os novos vetores de chave e valor são armazenados no cache e adicionados às chaves e valores passados. $$K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)$$

A atenção é calculada de forma independente em cada camada do modelo, e o cache é feito por camada.

Consulte a tabela abaixo para comparar como o cache melhora a eficiência.

Classe de cache

Uma interface básica de cache de KV recebe um tensor de chave e um tensor de valor para o token atual e retorna os tensors K e V atualizados. Isso é gerenciado internamente pelo método forward de um modelo.

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

Ao usar a classe Cache do Transformers, o módulo de autoatenção executa várias etapas críticas para integrar informações passadas e atuais.

1. O módulo de atenção concatena os pares kv atuais com os pares kv anteriores armazenados no cache. Isso cria pesos de atenção com a forma (new_tokens_length, past_kv_length + new_tokens_length). Os pares kv atuais e anteriores são essencialmente combinados para calcular as pontuações de atenção, garantindo que o modelo esteja ciente do contexto anterior e da entrada atual.

2. Quando o método forward é chamado iterativamente, é crucial que a forma da attention mask corresponda ao comprimento combinado dos pares kv passados e atuais. A attention mask deve ter a forma (batch_size, past_kv_length + new_tokens_length). Isso geralmente é tratado internamente em generate(), mas, se você quiser implementar seu próprio loop de geração com Cache, tenha isso em mente! A attention mask deve conter os valores dos tokens passados e atuais.

Implementação de armazenamento em cache

Os caches são estruturados como uma lista de camadas, em que cada camada contém um cache de chave e um cache de valor. Os caches de chave e de valor são tensores com a forma [batch_size, num_heads, seq_len, head_dim].

As camadas podem ser de diferentes tipos (por exemplo, DynamicLayer, StaticLayer, StaticSlidingWindowLayer), o que altera principalmente o modo como o comprimento da sequência é tratado e como o cache é atualizado.

A mais simples é uma DynamicLayer que cresce à medida que mais tokens são processados. A dimensão do comprimento da sequência (seq_len) aumenta a cada novo token:

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)

Outros tipos de camada, como StaticLayer e StaticSlidingWindowLayer, possuem um comprimento de sequência fixo que é definido quando o cache é criado. Isso os torna compatíveis com torch.compile. No caso de StaticSlidingWindowLayer, os tokens existentes são deslocados para fora do cache quando um novo token é adicionado.

O exemplo abaixo demonstra como criar um loop de geração com DynamicCache. Conforme discutido, a máscara de atenção é uma concatenação dos valores de tokens passados e atuais.

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,"