Modelo Llama 3 implementado em NumPy puro

• Entender com precisão a estrutura por meio de uma implementação realmente funcional do modelo Llama 3

Visão geral

• O modelo Llama 3, divulgado pela Meta, está chamando atenção.
• Ele se destaca por uma escala e desempenho impressionantes, com 24K GPUs, 15T de dados de treinamento, 10M de dados de instruções e 1.3M horas de GPU.
• A estrutura do modelo não mudou significativamente. O Llama 3 usa GQA, mas isso já havia sido implementado no Llama 2 70B.
• A implementação foi feita usando apenas NumPy para permitir uma compreensão intuitiva da estrutura do modelo.
• Foi utilizado o modelo stories15M, treinado por Andrej Karpathy com a arquitetura do Llama 2, convertido para o formato compactado do NumPy.

Estrutura

• A estrutura do modelo Llama 3 é a mesma do 42dot LLM.
• Parâmetros do modelo:
  • dim: 288
  • n_layers: 6
  • n_heads: 6
  • vocab_size: 32000
  • max_seq_len: 256
  • max_new_tokens: 50

RoPE #1

• Os valores de cos e sin são pré-calculados para o embedding RoPE.
• Esses valores são usados em Q e K.
• O resultado do cálculo é multiplicado com np.outer, e então cos e sin são calculados.

RMSNorm

• O RMSNorm normaliza os valores de ativação pela Root Mean Square, em vez de usar estatísticas tradicionais de Mini Batch ou Layer.
• Isso fornece um escalonamento consistente das ativações.

QKV

• Diferentemente do GPT, que calcula QKV fazendo um único matmul com os pesos e depois dividindo, o Llama possui pesos separados para cada um de Q, K e V.
• Cada valor é reorganizado para Multi-Head Attention.

RoPE #2

• O RoPE possui características tanto de codificação posicional absoluta quanto relativa.
• Ele é aplicado apenas a Q e K; a entrada é dividida, multiplicada por cos e sin, e depois recomposta por meio de soma e subtração dos resultados.

Cache KV

• Modelos generativos no estilo GPT usam Masked Attention, o que permite o uso de cache KV.
• Como os resultados anteriores são sempre os mesmos, K e V são armazenados em cache, e Q é calculado apenas para o último valor.

GQA (Grouped-Query Attention)

• GQA é uma técnica introduzida no Llama 2 que oferece economia de memória e melhora de desempenho.
• No Llama 3, o GQA é aplicado a todos os modelos acima de 8B.

Scaled Dot-Product Attention

• Cada atenção é calculada por meio de Multi-Head Attention.
• O resultado é obtido com softmax e matmul.

Feed Forward

• O Feed Forward do modelo Llama usa 3 camadas lineares e não tem bias.
• Ele gera o valor swish, multiplica por x_V e depois reduz a escala novamente.

SwiGLU

• O SwiGLU melhora o desempenho do modelo com uma combinação característica de várias camadas feed-forward.

Linear

• Na saída final, apenas o último logit é calculado com matmul para aumentar a velocidade.

Geração

• Os tokens são gerados um por um usando os logits extraídos.
• O processo é dividido em Prefill Phase e Decode Phase.
• Na Prefill Phase, toda a entrada é passada; na Decode Phase, apenas o ID do último token é passado para obter o resultado.

Exemplo

• É possível executar da seguinte forma:

  $ python llama3.py "I have a dream"  
  

GitHub

• O código-fonte completo pode ser visto em likejazz/llama3.np.

Referências

1. Exploring and Building the Llama 3 Architecture
2. Rotation Matrix
3. Mastering LLM Techniques: Inference Optimization
4. arXiv:2305.13245

Opinião do GN⁺

• Estrutura e desempenho do modelo Llama 3: O modelo Llama 3 mantém a estrutura do Llama 2 anterior, mas melhora bastante o desempenho. Isso é resultado de um projeto que considera ao mesmo tempo escalabilidade e eficiência.
• Por que implementar com NumPy: Ao implementar o modelo com NumPy, é possível entender de forma mais intuitiva sua estrutura e funcionamento. Isso ajuda bastante estudantes e pesquisadores.
• Introdução do GQA: O GQA é uma técnica que oferece ao mesmo tempo economia de memória e melhora de desempenho, e sua aplicação em todos os modelos do Llama 3 maximiza a eficiência do modelo.
• Importância do cache KV: O cache KV tem um papel importante em modelos generativos no estilo GPT, e por meio dele é possível aumentar bastante a eficiência computacional do modelo.
• Caso de uso prático: Com o código de exemplo, é possível executar o modelo na prática, o que é uma boa oportunidade para verificar diretamente seu desempenho.

[xguru]

O que foi postado no Hacker News estava em inglês, mas foi alterado para o link que o autor original, Likejazz, escreveu em coreano.