Otimizando LLMs para velocidade e memória

Grandes Modelos de Linguagem (LLMs), como GPT3/4, Falcon e Llama, estão avançando rapidamente em sua capacidade de realizar tarefas centradas no ser humano, consolidando-se como ferramentas essenciais em indústrias modernas baseadas em conhecimento. No entanto, a implantação desses modelos em tarefas do mundo real continua sendo desafiadora:

Para exibir capacidades de compreensão e geração de texto próximas às humanas, os LLMs atualmente precisam ser compostos por bilhões de parâmetros (veja Kaplan et al, Wei et. al). Isso, consequentemente, amplifica as exigências de memória para inferência.
Em muitas tarefas do mundo real, é necessário fornecer aos LLMs informações contextuais extensas. Isso exige que o modelo tenha a capacidade de gerenciar sequências de entrada muito longas durante a inferência.

O cerne desses desafios reside em ampliar as capacidades computacionais e de memória dos LLMs, especialmente ao lidar com sequências de entrada extensas.

Neste guia, vamos apresentar as técnicas eficazes para a implantação eficiente de LLMs:

Precisão reduzida: Pesquisas demonstraram que operar com precisão numérica reduzida, especificamente 8 bits e 4 bits, pode proporcionar vantagens computacionais sem uma queda considerável no desempenho do modelo.
Flash Attention: O Flash Attention é uma variação do algoritmo de atenção que não apenas oferece uma abordagem mais eficiente em termos de memória, mas também realiza um aumento na eficiência devido ao uso otimizado da memória da GPU.
Inovações Arquiteturais: Considerando que os LLMs são sempre implantados da mesma forma durante a inferência, a saber, geração de texto autoregressiva com um contexto de entrada longo, foram propostas arquiteturas de modelo especializadas que permitem uma inferência mais eficiente. Os avanços mais importantes nas arquiteturas de modelo são Alibi, Rotary embeddings, Multi-Query Attention (MQA) e Grouped-Query-Attention (GQA).

Ao longo deste guia, apresentaremos uma análise da geração auto-regressiva sob a perspectiva de tensores. Exploramos as vantagens e desvantagens da adoção de menor precisão, oferecemos uma exploração abrangente dos algoritmos de atenção mais recentes e discutimos arquiteturas de LLM aprimoradas. Enquanto fazemos isso, executamos exemplos práticos que destacam cada uma das melhorias de recursos.

1. Menor precisão

Os requisitos de memória dos LLMs podem ser melhor compreendidos ao se ver o LLM como um conjunto de matrizes e vetores de pesos e as entradas de texto como uma sequência de vetores. A seguir, a definição pesos será usada para significar todas as matrizes e vetores de pesos do modelo.

No momento em que este guia foi escrito, os LLMs consistem em pelo menos alguns bilhões de parâmetros. Cada parâmetro, portanto, é composto por um número decimal, por exemplo, 4.5689, que geralmente é armazenado no formato float32, bfloat16 ou float16. Isso nos permite calcular facilmente a memória necessária para carregar o LLM na memória:

Carregar os pesos de um modelo com X bilhões de parâmetros requer aproximadamente 4 * X GB de VRAM em precisão float32

Hoje em dia, no entanto, os modelos raramente são treinados em precisão float32 completa, mas geralmente em precisão bfloat16 ou, com menos frequência, em precisão float16. Portanto, a regra prática se torna:

Carregar os pesos de um modelo com X bilhões de parâmetros requer aproximadamente 2 * X GB de VRAM em precisão bfloat16/float16

Para entradas de texto mais curtas (menos de 1024 tokens), o requisito de memória para inferência é amplamente dominado pelo requisito de memória para carregar os pesos. Portanto, por enquanto, vamos assumir que o requisito de memória para inferência é igual ao requisito de memória para carregar o modelo na VRAM da GPU.

Para dar alguns exemplos de quanto VRAM é necessário, aproximadamente, para carregar um modelo em bfloat16:

GPT3 requer 2 * 175 GB = 350 GB de VRAM
Bloom requer 2 * 176 GB = 352 GB de VRAM
Llama-2-70b requer 2 * 70 GB = 140 GB de VRAM
Falcon-40b requer 2 * 40 GB = 80 GB de VRAM
MPT-30b requer 2 * 30 GB = 60 GB de VRAM
bigcode/starcoder requer 2 * 15,5 = 31 GB de VRAM

No momento da redação deste documento, o maior chip de GPU disponível no mercado é o A100 & H100, que oferece 80GB de VRAM. A maioria dos modelos listados anteriormente requer mais de 80GB apenas para ser carregada e, portanto, necessariamente requer tensor parallelism e/ou pipeline parallelism.

🤗 O Transformers agora oferece suporte a paralelismo de tensor para modelos compatíveis que possuem base_tp_plan em suas respectivas classes de configuração. Saiba mais sobre Paralelismo de Tensor aqui. Além disso, se você estiver interessado em escrever modelos de uma forma amigável ao paralelismo de tensor, sinta-se à vontade para dar uma olhada na biblioteca text-generation-inference.

O paralelismo de pipeline ingênuo é suportado nativamente. Para isso, basta carregar o modelo com device="auto", o que posiciona automaticamente as diferentes camadas nas GPUs disponíveis, conforme explicado aqui. Observe, no entanto, que, embora seja muito eficaz, esse paralelismo de pipeline ingênuo não resolve os problemas de ociosidade das GPUs. Para isso, é necessário um paralelismo de pipeline mais avançado, conforme explicado aqui.

Se você tiver acesso a um nó 8 x 80GB A100, poderá carregar o BLOOM da seguinte forma

!pip install transformers accelerate bitsandbytes optimum

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

Ao usar device_map="auto", as camadas de atenção seriam distribuídas igualmente entre todas as GPUs disponíveis.

Neste guia, usaremos bigcode/octocoder, pois ele pode ser executado em um único dispositivo de GPU A100 com 40 GB. Observe que todas as otimizações de memória e velocidade que aplicaremos a seguir são igualmente aplicáveis a modelos que exigem paralelismo de modelo ou de tensor.

Como o modelo é carregado em precisão bfloat16, usando nossa regra geral acima, esperaríamos que o requisito de memória para executar a inferência com bigcode/octocoder fosse de cerca de 31 GB de VRAM. Vamos tentar.

Primeiro, carregamos o modelo e o tokenizer e, em seguida, passamos ambos para o objeto pipeline do Transformers.

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Saída:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

Ótimo, agora podemos usar diretamente o resultado para converter bytes em Gigabytes.

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

Vamos chamar torch.accelerator.memory.max_memory_allocated para medir a alocação de memória de pico do acelerador.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Saída:

29.0260648727417

Bastante próximo do nosso cálculo de estimativa! Podemos ver que o número não está exatamente correto, pois a conversão de bytes para kilobytes exige uma multiplicação por 1024, e não por 1000. Portanto, a fórmula de estimativa também pode ser entendida como um cálculo de “no máximo X GB”. Observe que, se tivéssemos tentado executar o modelo em precisão float32 completa, teríamos precisado de nada menos que 64 GB de VRAM.

Quase todos os modelos são treinados em bfloat16 atualmente, não há motivo para executar o modelo em precisão completa float32 se sua GPU suporta bfloat16. Float32 não fornecerá resultados de inferência melhores do que a precisão usada para treinar o modelo.

Se você não tiver certeza sobre o formato em que os pesos do modelo estão armazenados no Hub, você sempre pode verificar a configuração do checkpoint em "dtype", por exemplo aqui. É recomendável definir o modelo com o mesmo tipo de precisão indicado na configuração ao carregar com from_pretrained(..., dtype=...), exceto quando o tipo original é float32, caso em que é possível usar tanto float16 quanto bfloat16 para inferência.

Vamos definir uma função flush(...) para liberar toda a memória alocada, de modo que possamos medir com precisão a memória de GPU alocada no pico.

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.accelerator.empty_cache()
  torch.accelerator.reset_peak_memory_stats()

Vamos chamá-lo agora para o próximo experimento.

flush()

A partir da biblioteca Accelerate, você também pode usar um método de utilidade agnóstico de dispositivo chamado release_memory, que leva em consideração diversos backends de hardware, como XPU, MLU, NPU, MPS e outros.

from accelerate.utils import release_memory
# ...

release_memory(model)

E se a sua GPU não tiver 32 GB de VRAM? Foi constatado que os pesos do modelo podem ser quantizados para 8 bits ou 4 bits sem uma perda significativa de desempenho (veja Dettmers et al.). O modelo pode ser quantizado até para 3 ou 2 bits com uma perda de desempenho aceitável, conforme mostrado no recente artigo GPTQ 🤯.

Sem entrar em muitos detalhes, os esquemas de quantização visam reduzir a precisão dos pesos, tentando manter os resultados de inferência do modelo o mais preciso possível (ou seja, o mais próximo possível do bfloat16). Observe que a quantização funciona especialmente bem para geração de texto, pois tudo o que nos importa é escolher o conjunto de tokens mais prováveis a seguir e não nos importamos muito com os valores exatos da distribuição de logit do próximo token. O que importa é que a distribuição de logit do próximo token permaneça aproximadamente a mesma, para que uma operação de argmax ou topk produza os mesmos resultados.

Existem diversas técnicas de quantização, que não detalharemos aqui, mas, em geral, todas as técnicas de quantização funcionam da seguinte forma:

1. Quantize todos os pesos para a precisão de destino
1. Carregue os pesos quantizados e passe a sequência de vetores de entrada em precisão bfloat16
1. Desquantize dinamicamente os pesos para bfloat16 para executar o cálculo com seus vetores de entrada em precisão bfloat16

Em resumo, isso significa que multiplicações da matriz de pesos pelas entradas com $X$ sendo as entradas, $W$ sendo uma matriz de pesos e $Y$ sendo a saída: $Y = X * W$

são alteradas para $Y = X * \text{dequantize}(W)$

para cada multiplicação de matrizes. A desquantização e a requantização são executadas sequencialmente para todas as matrizes de pesos conforme as entradas percorrem o grafo da rede.

Portanto, o tempo de inferência muitas vezes não é reduzido ao usar pesos quantizados, mas sim aumentado. Chega de teoria, vamos tentar! Para quantizar os pesos com Transformers, você precisa garantir que a biblioteca bitsandbytes esteja instalada.

!pip install bitsandbytes

Em seguida, podemos carregar modelos em quantização de 8 bits simplesmente adicionando a flag load_in_8bit=True ao from_pretrained.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_8bit=True), pad_token_id=0)

Agora, vamos executar nosso exemplo novamente e medir o uso de memória.

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Saída:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

Ótimo, estamos obtendo o mesmo resultado de antes, então não houve perda de precisão! Vamos ver quanto de memória foi usada desta vez.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Saída:

15.219234466552734

Muito menos! Estamos com apenas um pouco mais de 15 GBs e, portanto, podemos executar este modelo em GPUs de consumo, como a 4090. Estamos observando um ganho muito bom na eficiência de memória e praticamente nenhuma degradação na saída do modelo. No entanto, também podemos notar uma leve redução na velocidade durante a inferência.

Excluímos os modelos e liberamos a memória novamente.

del model
del pipe

flush()

Vamos ver qual é o consumo de memória de GPU no pico que a quantização de 4 bits proporciona. A quantização do modelo para 4 bits pode ser feita com a mesma API de antes — desta vez, passando load_in_4bit=True em vez de load_in_8bit=True.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_4bit=True), pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Saída:

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

Estamos quase vendo o mesmo texto de saída de antes - apenas o python está faltando logo antes do trecho de código. Vamos ver quanta memória foi necessária.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Saída:

9.543574333190918

Apenas 9,5 GB! Isso realmente não é muito para um modelo com mais de 15 bilhões de parâmetros.

Embora observemos uma degradação muito pequena na precisão do nosso modelo aqui, a quantização de 4 bits pode, na prática, frequentemente levar a resultados diferentes em comparação com a quantização de 8 bits ou a inferência completa em bfloat16. Cabe ao usuário testá-la.

Observe também que a inferência aqui foi novamente um pouco mais lenta em comparação com a quantização de 8 bits, o que se deve ao método de quantização mais agressivo usado para a quantização de 4 bits, resultando em $\text{quantize}$ e $\text{dequantize}$ levando mais tempo durante a inferência.

del model
del pipe

flush()

No geral, vimos que executar o OctoCoder em precisão de 8 bits reduziu a VRAM de GPU necessária de 32 GB para apenas 15 GB, e executar o modelo em precisão de 4 bits reduziu ainda mais a VRAM de GPU necessária para pouco mais de 9 GB.

A quantização de 4 bits permite que o modelo seja executado em GPUs como RTX3090, V100 e T4, que são bastante acessíveis para a maioria das pessoas.

Para mais informações sobre quantização e para ver como é possível quantizar modelos para exigir ainda menos memória de VRAM da GPU do que 4 bits, recomendamos dar uma olhada na implementação GPT-QModel.

Como conclusão, é importante lembrar que a quantização de modelos troca a eficiência de memória aprimorada por precisão e, em alguns casos, pelo tempo de inferência.

Se a memória da GPU não for uma restrição para o seu caso de uso, muitas vezes não há necessidade de investigar a quantização. No entanto, muitas GPUs simplesmente não conseguem executar LLMs sem métodos de quantização, e, nesse caso, os esquemas de quantização de 4 bits e 8 bits são ferramentas extremamente úteis.

Para informações de uso mais detalhadas, recomendamos fortemente consultar a Documentação de Quantização do Transformers. A seguir, vamos examinar como podemos melhorar a eficiência computacional e de memória usando algoritmos melhores e uma arquitetura de modelo aprimorada.

2. Flash Attention

Os LLMs de melhor desempenho atuais compartilham, mais ou menos, a mesma arquitetura fundamental, que consiste em camadas feed-forward, camadas de ativação, camadas de normalização de camada e, mais crucialmente, camadas de autoatenção.

As camadas de autoatenção são centrais nos Large Language Models (LLMs), pois permitem que o modelo compreenda as relações contextuais entre os tokens de entrada. No entanto, o consumo máximo de memória de GPU para as camadas de autoatenção cresce de forma quadrática tanto em complexidade de computação quanto de memória em função do número de tokens de entrada (também chamado de comprimento da sequência) , que denotamos a seguir por $N$ . Embora isso não seja muito perceptível para sequências de entrada mais curtas (de até 1000 tokens de entrada), torna-se um problema sério para sequências de entrada mais longas (em torno de 16000 tokens de entrada).

Vamos dar uma olhada mais de perto. A fórmula para calcular a saída $\mathbf{O}$ de uma camada de autoatenção para uma entrada $\mathbf{X}$ de comprimento $N$ é: $\textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X}$

$\mathbf{X} = (\mathbf{x}_1, ... \mathbf{x}_{N})$ é, portanto, a sequência de entrada para a camada de atenção. As projeções $\mathbf{Q}$ e $\mathbf{K}$ serão compostas, cada uma, por $N$ vetores, resultando em $\mathbf{QK}^T$ de tamanho $N^2$ .

Os LLMs geralmente possuem múltiplas cabeças de atenção, realizando assim vários cálculos de autoatenção em paralelo. Supondo que o LLM tenha 40 cabeças de atenção e execute em precisão bfloat16, podemos calcular a memória necessária para armazenar as $\mathbf{QK^T}$ matrizes como $40 * 2 * N^2$ bytes. Para $N=1000$ são necessários apenas cerca de 50 MB de VRAM, no entanto, para $N=16000$ seriam necessários 19 GB de VRAM, e para $N=100,000$ seriam necessários quase 1 TB apenas para armazenar as $\mathbf{QK}^T$ matrizes.

Em resumo, o algoritmo padrão de autoatenção rapidamente se torna proibitivamente caro em termos de memória para contextos de entrada grandes.

À medida que as LLMs aprimoram a compreensão e a geração de texto, são aplicadas a tarefas cada vez mais complexas. Enquanto os modelos antes lidavam com a tradução ou a sumarização de algumas frases, agora gerenciam páginas inteiras, exigindo a capacidade de processar extensos comprimentos de entrada.

Como podemos eliminar os requisitos de memória exorbitantes para grandes comprimentos de entrada? Precisamos de uma nova maneira de calcular o mecanismo de autoatenção que elimine a $\mathbf{QK}^T$ matriz. Tri Dao et al. desenvolveram exatamente esse novo algoritmo e o chamaram de Flash Attention.

Em resumo, o Flash Attention separa o cálculo de $\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T)$ e, em vez disso, calcula blocos menores da saída iterando por meio de múltiplas etapas de cálculo de softmax: $\textbf{O}_i \leftarrow s^a_{ij} * \textbf{O}_i + s^b_{ij} * \mathbf{V}_{j} \times \text{Softmax}(\mathbf{QK}^T_{i,j}) \text{ for multiple } i, j \text{ iterations}$

com $s^a_{ij}$ e $s^b_{ij}$ sendo algumas estatísticas de normalização softmax que precisam ser recalculadas para cada $i$ e $j$ .

Observe que o Flash Attention como um todo é um pouco mais complexo e está bastante simplificado aqui, pois aprofundar-se demais está fora do escopo deste guia. O leitor é convidado a consultar o bem redigido artigo do Flash Attention para mais detalhes.

O principal ponto a ser destacado aqui é:

Ao manter o controle das estatísticas de normalização do softmax e ao usar algumas matemáticas inteligentes, o Flash Attention fornece saídas numericamente idênticas em comparação com a camada de autoatenção padrão, com um custo de memória que aumenta apenas linearmente com $N$ .

Ao observar a fórmula, diria-se intuitivamente que o Flash Attention deve ser muito mais lento em comparação com a fórmula padrão de autoatenção, pois é necessário realizar mais cálculos. De fato, o Flash Attention requer mais FLOPs em comparação com a atenção normal, pois as estatísticas de normalização do softmax precisam ser constantemente recalculadas (veja o paper para mais detalhes, se interessado)

No entanto, o Flash Attention é muito mais rápido na inferência em comparação com a atenção padrão, devido à sua capacidade de reduzir significativamente as demandas sobre a memória de alta largura de banda, mais lenta, da GPU (VRAM), focando em vez disso na memória on-chip, mais rápida (SRAM).

Em essência, o Flash Attention garante que todas as operações intermediárias de escrita e leitura possam ser realizadas usando a memória on-chip SRAM, em vez de precisar acessar a memória VRAM mais lenta para calcular o vetor de saída $\mathbf{O}$ .

Na prática, atualmente não há absolutamente nenhum motivo para não usar Flash Attention, caso esteja disponível. O algoritmo produz saídas matematicamente idênticas e é tanto mais rápido quanto mais eficiente em termos de memória.

3. Inovações Arquiteturais

Até agora, analisamos a melhoria da eficiência computacional e de memória por meio de:

Casting dos pesos para um formato de menor precisão
Substituindo o algoritmo de autoatenção por uma versão mais eficiente em termos de memória e computação

Vamos agora examinar como podemos alterar a arquitetura de um LLM para que ele seja mais eficaz e eficiente em tarefas que exigem entradas de texto longas, por exemplo:

Respostas a perguntas com recuperação aumentada,
Resumo,
Chat

Observe que chat não apenas exige que o LLM processe entradas de texto longas, mas também requer que o LLM seja capaz de lidar de forma eficiente com o diálogo de ida e volta entre usuário e assistente (como o ChatGPT).

Uma vez treinado, a arquitetura fundamental do LLM é difícil de alterar, portanto, é importante fazer considerações sobre as tarefas do LLM previamente e, de acordo com isso, otimizar a arquitetura do modelo. Existem dois componentes importantes da arquitetura do modelo que rapidamente se tornam gargalos de memória e/ou desempenho para sequências de entrada grandes.

Os embeddings posicionais
O cache de chave-valor

Vamos analisar cada componente em mais detalhes

3.1 Melhorando os embeddings posicionais de LLMs

A autoatenção coloca cada token em relação aos tokens dos demais. Por exemplo, a $\text{Softmax}(\mathbf{QK}^T)$ matriz da sequência de entrada de texto “Hello”, “I”, “love”, “you” poderia ser a seguinte:

Cada token de palavra recebe uma massa de probabilidade com a qual atende a todos os outros tokens de palavra e, portanto, é colocado em relação com todos os outros tokens de palavra. Por exemplo, a palavra “love” atende à palavra “Hello” com 5%, a “I” com 30% e a si mesma com 65%.

Uma LLM baseada em autoatenção, mas sem embeddings de posição, teria grandes dificuldades em compreender as posições relativas dos tokens de entrada. Isso ocorre porque a pontuação de probabilidade calculada por $\mathbf{QK}^T$ relaciona cada token de palavra a todos os outros tokens em $O(1)$ computações, independentemente da distância posicional relativa entre eles. Portanto, em uma LLM sem embeddings de posição, cada token parece estar à mesma distância de todos os outros, e.g. diferenciar entre “Hello I love you e “You love I hello seria muito desafiador.

Para que o LLM compreenda a ordem das frases, é necessário um sinal adicional, geralmente aplicado na forma de codificações posicionais (também chamadas de embeddings posicionais). As codificações posicionais codificam a posição de cada token em uma representação numérica que o LLM pode aproveitar para compreender melhor a ordem das frases.

Os autores do Attention Is All You Need introduziram embeddings posicionais sinusoidais $\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N$ . em que cada vetor $\mathbf{p}_i$ é calculado como uma função sinusoidal de sua posição $i$ . Os codificadores posicionais são então simplesmente somados aos vetores da sequência de entrada $\mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N$ = $\mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N$ indicando ao modelo a ordem das frases para uma melhor aprendizagem.

Em vez de usar embeddings de posição fixos, outros (como Devlin et al.) usaram codificações posicionais aprendidas, para as quais os embeddings de posição $\mathbf{P}$ são aprendidos durante o treinamento.

Os embeddings de posição sinusoidais e aprendidos costumavam ser os métodos predominantes para codificar a ordem das frases em LLMs, mas foram encontrados alguns problemas relacionados a essas codificações posicionais:

Os embeddings posicionais sinusoidais e aprendidos são ambos embeddings posicionais absolutos, i.e. ou seja, codificam um embedding único para cada id de posição: $0, \ldots, N$ . Como demonstrado por Huang et al. e Su et al., os embeddings posicionais absolutos resultam em baixo desempenho de LLMs para entradas de texto longas. Para entradas de texto longas, é vantajoso que o modelo aprenda a distância posicional relativa entre os tokens de entrada, em vez de sua posição absoluta.
Ao usar embeddings de posição aprendidos, o LLM precisa ser treinado em um comprimento de entrada fixo $N$, o que dificulta a extrapolação para um comprimento de entrada maior do que aquele em que foi treinado.

Recentemente, incorporações posicionais relativas que podem lidar com os problemas mencionados acima tornaram-se mais populares, em especial:

Ambos RoPE e ALiBi argumentam que o ideal é indicar a ordem das frases ao LLM diretamente no algoritmo de autoatenção, pois é nele que os tokens de palavras são colocados em relação uns com os outros. Mais especificamente, a ordem das frases deve ser indicada modificando o $\mathbf{QK}^T$ cálculo.

Sem entrar em muitos detalhes, RoPE observa que as informações de posição podem ser codificadas em pares de consulta-chave, e.g. $\mathbf{q}_i$ e $\mathbf{x}_j$ girando cada vetor por um ângulo $\theta * i$ e $\theta * j$ respectivamente, com $i, j$ descrevendo a posição da frase de cada vetor: $\mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j.$

$\mathbf{R}_{\theta, i - j}$ representa, portanto, uma matriz de rotação. $\theta$ é não aprendido durante o treinamento, mas definido com um valor pré-definido que depende do comprimento máximo da sequência de entrada durante o treinamento.

Dessa forma, a pontuação de probabilidade entre $\mathbf{q}_i$ e $\mathbf{q}_j$ é afetada apenas se $i \ne j$ e depende exclusivamente da distância relativa $i - j$ independentemente das posições específicas de cada vetor $i$ e $j$ .

RoPE é usado em vários dos LLMs mais importantes de hoje, como:

Como alternativa, ALiBi propõe um esquema de codificação de posição relativa muito mais simples. A distância relativa entre os tokens de entrada é adicionada como um inteiro negativo, escalado por um valor predefinido m a cada entrada de consulta-chave da $\mathbf{QK}^T$ matriz, logo antes do cálculo do softmax.

Como mostrado no paper ALiBi, essa codificação posicional relativa simples permite que o modelo mantenha um alto desempenho mesmo em sequências de entrada de texto muito longas.

ALiBi é usado em vários dos LLMs mais importantes de hoje, como:

Ambos RoPE e ALiBi posições de codificação podem extrapolar para comprimentos de entrada não vistos durante o treinamento, enquanto foi demonstrado que a extrapolação funciona muito melhor nativamente para ALiBi em comparação com RoPE. Para ALiBi, basta aumentar os valores da matriz de posição triangular inferior para corresponder ao comprimento da sequência de entrada. Para RoPE, manter o mesmo $\theta$ que foi usado durante o treinamento leva a resultados ruins ao passar entradas de texto muito mais longas do que as vistas durante o treinamento, c.f Press et al.. No entanto, a comunidade encontrou alguns truques eficazes que adaptam $\theta$, permitindo assim que RoPE embeddings de posição funcionem bem para sequências de entrada de texto extrapoladas (veja aqui).

Tanto o RoPE quanto o ALiBi são incorporações posicionais relativas que não são aprendidas durante o treinamento, mas sim baseadas nas seguintes intuições:

As pistas posicionais sobre as entradas de texto devem ser fornecidas diretamente à $\mathbf{QK}^T$ matriz da camada de autoatenção.
O LLM deve ser incentivado a aprender uma codificação posicional de distância relativa constante.
Quanto mais distantes entre si estão os tokens de entrada de texto, menor é a probabilidade de consulta-valor deles. Tanto o RoPE quanto o ALiBi reduzem a probabilidade de consulta-chave de tokens distantes entre si. O RoPE reduz essa probabilidade diminuindo o produto vetorial ao aumentar o ângulo entre os vetores de consulta-chave. O ALiBi reduz ao somar números negativos grandes ao produto vetorial.

Em conclusão, LLMs destinados a tarefas que exigem o processamento de entradas de texto grandes são melhor treinados com embeddings posicionais relativos, como RoPE e ALiBi. Observe também que, mesmo que um LLM com RoPE e ALiBi tenha sido treinado apenas em um comprimento fixo, por exemplo, $N_1 = 2048$ ele ainda pode ser usado na prática com entradas de texto muito maiores que $N_1$, como $N_2 = 8192 > N_1$, por meio da extrapolação dos embeddings posicionais.

3.2 O cache de chave-valor

A geração de texto auto-regressiva com LLMs funciona inserindo iterativamente uma sequência de entrada, amostrando o próximo token, adicionando esse token à sequência de entrada e continuando esse processo até que o LLM produza um token que indique o término da geração.

Dê uma olhada no Tutorial de Geração de Texto do Transformers para obter uma explicação mais visual de como funciona a geração auto-regressiva.

Vamos executar um trecho de código rápido para mostrar como a geração auto-regressiva funciona na prática. Simplesmente vamos obter o próximo token mais provável por meio de torch.argmax.

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("shape of input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

Saída:

shape of input_ids torch.Size([1, 21])
shape of input_ids torch.Size([1, 22])
shape of input_ids torch.Size([1, 23])
shape of input_ids torch.Size([1, 24])
shape of input_ids torch.Size([1, 25])
[' Here is a Python function']

Como podemos ver, a cada vez que aumentamos os tokens de entrada do texto pelo token recém-amostrado.

Com exceção de poucos casos, os LLMs são treinados usando o objetivo de modelagem de linguagem causal e, portanto, mascaram o triângulo superior da matriz de pontuações de atenção - é por isso que, nos dois diagramas acima, as pontuações de atenção ficam em branco (ou seja, têm probabilidade 0). Para uma breve revisão sobre a modelagem de linguagem causal, você pode consultar o blog Illustrated Self Attention.

Como consequência, os tokens nunca dependem de tokens posteriores, mais especificamente o vetor $\mathbf{q}_i$ nunca é colocado em relação com vetores de chave ou valor $\mathbf{k}_j, \mathbf{v}_j$ se $j > i$. Em vez disso, $\mathbf{q}_i$ apenas presta atenção aos vetores de chave-valor anteriores $\mathbf{k}{m < i}, \mathbf{v}{m < i} \text{ , para } m \in {0, \ldots i - 1}$. Para reduzir cálculos desnecessários, é possível, portanto, armazenar em cache os vetores de chave-valor de cada camada para todos os passos de tempo anteriores.

A seguir, instruiremos o LLM a utilizar o cache de chave-valor, recuperando-o e encaminhando-o a cada passagem de frente. No Transformers, podemos recuperar o cache de chave-valor passando a flag use_cache para a chamada forward e, em seguida, passá-lo junto com o token atual.

past_key_values = None # past_key_values is the key-value cache
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", next_token_id.shape)
  print("length of key-value cache", past_key_values.get_seq_length())  # past_key_values are of shape [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

Saída:

shape of input_ids torch.Size([1, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
length of key-value cache 24
[' Here', ' is', ' a', ' Python', ' function']

Como se pode observar, ao usar o cache de chave-valor, os tokens de entrada de texto não têm seu comprimento aumentado, mas permanecem como um único vetor de entrada. O comprimento do cache de chave-valor, por outro lado, é aumentado em um a cada etapa de decodificação.

Ao fazer uso do cache de chave-valor, o $\mathbf{QK}^T$ é essencialmente reduzido a $\mathbf{q}_c\mathbf{K}^T$ com $\mathbf{q}_c$ sendo a projeção de consulta do token de entrada atualmente processado, que é sempre apenas um único vetor.

O uso do cache de chave-valor tem duas vantagens:

Aumento significativo na eficiência computacional, pois são realizadas menos operações em comparação com o cálculo da matriz completa $\mathbf{QK}^T$ matriz. Isso resulta em um aumento na velocidade de inferência
A memória máxima necessária não aumenta de forma quadrática com o número de tokens gerados, mas apenas de forma linear.

Deve-se sempre utilizar o cache de chave-valor, pois ele leva a resultados idênticos e a uma aceleração significativa para sequências de entrada mais longas. O Transformers tem o cache de chave-valor habilitado por padrão ao utilizar o pipeline de texto ou o método generate. Temos um guia inteiro dedicado a caches aqui.

Observe que, apesar de nossa recomendação de usar caches de chave-valor, a saída do seu LLM pode ser ligeiramente diferente ao utilizá-los. Isso é uma propriedade dos próprios kernels de multiplicação de matrizes — você pode ler mais sobre isso aqui.

3.2.1 Conversa em várias rodadas

O cache de chave-valor é especialmente útil para aplicações, como chats, que exigem múltiplas passagens de decodificação auto-regressiva. Vamos examinar um exemplo.

User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants

Neste chat, o LLM executa a decodificação auto-regressiva duas vezes:

Na primeira vez, o cache de chave-valor está vazio e o prompt de entrada é "User: How many people live in France?" e o modelo gera o texto "Roughly 75 million people live in France" de forma auto-regressiva, aumentando o cache de chave-valor a cada etapa de decodificação.
Na segunda vez, o prompt de entrada é "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?". Graças ao cache, todos os vetores chave-valor das duas primeiras frases já foram calculados. Portanto, o prompt de entrada consiste apenas em "User: And how many in Germany?". Durante o processamento do prompt de entrada encurtado, seus vetores chave-valor calculados são concatenados ao cache de chave-valor da primeira decodificação. A resposta do segundo Assistant, "Germany has ca. 81 million inhabitants", é então gerada de forma auto-regressiva com o cache de chave-valor consistindo dos vetores chave-valor codificados de "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?".

Dois pontos merecem destaque aqui:

Manter todo o contexto é crucial para LLMs implantados em chat, para que o LLM compreenda todo o contexto anterior da conversa. Por exemplo, no exemplo acima, o LLM precisa entender que o usuário se refere à população ao perguntar "And how many are in Germany".
O cache de chave-valor é extremamente útil para chat, pois permite que o histórico de chat codificado cresça continuamente, em vez de ter que re-codificar o histórico de chat do zero (como, por exemplo, ocorreria ao usar uma arquitetura encoder-decoder).

No transformers, uma chamada de generate retornará past_key_values quando return_dict_in_generate=True for passado, além do padrão use_cache=True. Observe que isso ainda não está disponível por meio da interface pipeline.

# Generation as usual
prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True)
decoded_output = tokenizer.batch_decode(generation_output.sequences)[0]

# Piping the returned `past_key_values` to speed up the next conversation round
prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(
  **model_inputs,
  past_key_values=generation_output.past_key_values,
  max_new_tokens=60,
  return_dict_in_generate=True
)
tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

Saída:

 is a modified version of the function that returns Mega bytes instead.

def bytes_to_megabytes(bytes):
   return bytes / 1024 / 1024

Answer: The function takes a number of bytes as input and returns the number of

Ótimo, não há tempo adicional gasto para recalcular as mesmas chaves e valores para a camada de atenção! No entanto, há uma ressalva. Embora a memória de pico necessária para a $\mathbf{QK}^T$ matriz seja significativamente reduzida, manter o cache de chave-valor em memória pode se tornar muito custoso em termos de memória para sequências de entrada longas ou chats de múltiplas rodadas. Lembre-se de que o cache de chave-valor precisa armazenar os vetores de chave-valor para todos os vetores de entrada anteriores $\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$ em todas as camadas de autoatenção e para todas as cabeças de atenção.

Vamos calcular o número de valores float que precisam ser armazenados no cache de chave-valor para o LLM bigcode/octocoder que usamos antes. O número de valores float equivale a duas vezes o comprimento da sequência vezes o número de cabeças de atenção vezes a dimensão da cabeça de atenção e vezes o número de camadas. Calculando isso para nosso LLM em um comprimento hipotético de sequência de entrada de 16000, obtemos:

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

Saída:

7864320000

Aproximadamente 8 bilhões de valores float! Armazenar 8 bilhões de valores float em precisão float16 requer cerca de 15 GB de RAM, o que é aproximadamente metade do tamanho dos próprios pesos do modelo! Pesquisadores propuseram dois métodos que permitem reduzir significativamente o custo de memória do armazenamento do cache de chave-valor, que são explorados nas próximas subseções.

3.2.2 Multi-Query-Attention (MQA)

Multi-Query-Attention foi proposto no artigo Fast Transformer Decoding: One Write-Head is All You Need, de Noam Shazeer. Como o título indica, Noam descobriu que, em vez de usar pesos de projeção chave-valor para n_head, é possível usar um único par de pesos de projeção chave-valor compartilhado entre todas as cabeças de atenção, sem que o desempenho do modelo se degrade significativamente.

Ao usar um único par de pesos de projeção chave-valor, os vetores de chave e valor $\mathbf{k}_i, \mathbf{v}_i$ devem ser idênticos em todas as cabeças de atenção, o que significa que precisamos armazenar apenas 1 par de projeção chave-valor no cache, em vez de n_head pares.

Como a maioria dos LLMs usa entre 20 e 100 cabeças de atenção, o MQA reduz significativamente o consumo de memória do cache de chave-valor. Para o LLM usado neste notebook, foi possível, portanto, reduzir o consumo de memória necessário de 15 GB para menos de 400 MB em um comprimento de sequência de entrada de 16000.

Além da economia de memória, o MQA também resulta em maior eficiência computacional, conforme explicado a seguir. Na decodificação auto-regressiva, grandes vetores de chave-valor precisam ser recarregados e concatenados com o par de vetores de chave-valor atual para serem então inseridos no $\mathbf{q}_c\mathbf{K}^T$ cálculo a cada etapa. Na decodificação auto-regressiva, a largura de banda de memória necessária para o recarregamento constante pode se tornar um gargalo de tempo significativo. Ao reduzir o tamanho dos vetores de chave-valor, menos memória precisa ser acessada, reduzindo assim o gargalo de largura de banda de memória. Para mais detalhes, consulte o artigo de Noam.

A parte importante a ser compreendida aqui é que reduzir o número de cabeças de atenção chave-valor para 1 só faz sentido se um cache de chave-valor for utilizado. O consumo de memória de pico do modelo para uma única passagem de frente sem cache de chave-valor permanece inalterado, pois cada cabeça de atenção ainda possui um vetor de consulta único, de modo que cada cabeça de atenção ainda possui uma matriz $\mathbf{QK}^T$ T

O MQA tem sido amplamente adotado pela comunidade e é atualmente utilizado por muitos dos LLMs mais populares:

Além disso, o checkpoint usado neste notebook - bigcode/octocoder - faz uso de MQA.

3.2.3 Grouped-Query-Attention (GQA)

Grouped-Query-Attention, proposta por Ainslie et al. do Google, constatou que o uso de MQA pode frequentemente levar a uma degradação de qualidade em comparação com o uso de projeções multi-chave-valor padrão. O artigo argumenta que mais desempenho do modelo pode ser mantido reduzindo de forma menos drástica o número de pesos de projeção de cabeçalhos de consulta. Em vez de usar apenas um único peso de projeção chave-valor, devem ser usados n < n_head pesos de projeção chave-valor. Ao escolher n para um valor significativamente menor que n_head, como 2, 4 ou 8, quase todos os ganhos de memória e velocidade do MQA podem ser mantidos, sacrificando menos capacidade do modelo e, portanto, presumivelmente menos desempenho.

Além disso, os autores do GQA descobriram que checkpoints de modelos existentes podem ser uptrained para ter uma arquitetura GQA com apenas 5% do custo computacional do pré-treino original. Embora 5% do custo computacional do pré-treino original ainda possa ser uma quantidade enorme, o uptraining do GQA permite que checkpoints existentes sejam úteis para sequências de entrada mais longas.

O GQA foi proposto apenas recentemente, motivo pelo qual sua adoção ainda é menor no momento da escrita deste notebook. A aplicação mais notável do GQA é o Llama-v2.

Como conclusão, é fortemente recomendável utilizar GQA ou MQA quando o LLM for implantado com decodificação auto-regressiva e precisar lidar com sequências de entrada grandes, como é o caso, por exemplo, de chat.

Conclusão

A comunidade de pesquisa está constantemente criando novas e engenhosas maneiras de acelerar o tempo de inferência para LLMs cada vez maiores. Como exemplo, uma dessas direções de pesquisa promissoras é o decodificação especulativa, em que os “tokens fáceis” são gerados por modelos de linguagem menores e mais rápidos, e apenas os “tokens difíceis” são gerados pelo próprio LLM. Entrar em mais detalhes está fora do escopo deste notebook, mas pode ser lido neste bom post de blog.

A razão pela qual LLMs massivos, como GPT3/4, Llama-2-70b, Claude e PaLM, podem ser executados tão rapidamente em interfaces de chat como Hugging Face Chat ou ChatGPT se deve, em grande parte, às melhorias mencionadas acima em precisão, algoritmos e arquitetura. No futuro, aceleradores como GPUs, TPUs, etc… ficarão apenas mais rápidos e permitirão mais memória, mas ainda assim é importante sempre garantir o uso dos melhores algoritmos e arquiteturas disponíveis para obter o máximo retorno pelo investimento 🤗

Update on GitHub