Modelo Llama 3 implementado con NumPy puro

• Entender con precisión la arquitectura mediante una implementación realmente funcional del modelo Llama 3

Resumen

• El modelo Llama 3 publicado por Meta está llamando la atención.
• Presume una escala y un rendimiento abrumadores: 24K GPUs, 15T de datos de entrenamiento, 10M de datos de instrucciones y 1.3M de horas de GPU.
• La arquitectura del modelo no ha cambiado de forma drástica. Llama 3 usa GQA, pero esto ya se había implementado en Llama 2 70B.
• Está implementado usando solo NumPy para que la arquitectura del modelo pueda entenderse de forma intuitiva.
• Se usa el modelo stories15M, entrenado por Andrej Karpathy con la arquitectura de Llama 2, convertido al formato comprimido de NumPy.

Estructura

• La arquitectura del modelo Llama 3 es igual a la de 42dot LLM.
• Parámetros del modelo:
  • dim: 288
  • n_layers: 6
  • n_heads: 6
  • vocab_size: 32000
  • max_seq_len: 256
  • max_new_tokens: 50

RoPE #1

• Se precalculan cos y sin para la incrustación RoPE.
• Estos valores se usan en Q y K.
• El resultado del cálculo se multiplica con np.outer y luego se calculan cos y sin.

RMSNorm

• RMSNorm normaliza los valores de activación con Root Mean Square en lugar de usar estadísticas tradicionales de Mini Batch o Layer.
• Proporciona un escalado consistente de las activaciones.

QKV

• A diferencia de GPT, donde el cálculo de QKV se hace aplicando matmul a un solo peso y luego dividiéndolo, Llama tiene pesos separados para cada uno de QKV.
• Cada valor se reorganiza para Multi-Head Attention.

RoPE #2

• RoPE tiene propiedades tanto de codificación posicional absoluta como relativa.
• Se aplica solo a Q y K; divide la entrada, la multiplica por cos y sin, luego suma y resta los resultados para reconstruirla.

Caché KV

• Los modelos generativos estilo GPT usan Masked Attention, por lo que es posible usar caché KV.
• Como los resultados previos siempre son iguales, se almacenan K y V en caché, y Q se calcula solo para el último valor.

GQA (Grouped-Query Attention)

• GQA es una técnica introducida en Llama 2 que ofrece ahorro de memoria y mejora del rendimiento.
• En Llama 3, GQA se aplica a todos los modelos de 8B o más.

Scaled Dot-Product Attention

• Cada attention se calcula con Multi-Head Attention.
• El resultado se obtiene con softmax y matmul.

Feed Forward

• El Feed Forward del modelo Llama usa 3 capas lineales y no tiene bias.
• Genera el valor swish, lo multiplica por x_V y luego vuelve a reducir la escala.

SwiGLU

• SwiGLU mejora el rendimiento del modelo mediante una combinación distintiva de varias capas feed-forward.

Linear

• La salida final calcula solo el último logit con matmul para acelerar el proceso.

Generación

• Usa los logits extraídos para generar tokens uno por uno.
• Se divide en Prefill Phase y Decode Phase.
• En la Prefill Phase se pasa toda la entrada, y en la Decode Phase se pasa solo el ID del último token para obtener el resultado.

Ejemplo

• Puede ejecutarse así:

  $ python llama3.py "I have a dream"  
  

GitHub

• El código fuente completo puede revisarse en likejazz/llama3.np.

Referencias

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

Opinión de GN⁺

• La arquitectura y el rendimiento del modelo Llama 3: el modelo Llama 3 mantiene la arquitectura de Llama 2, pero mejora significativamente el rendimiento. Esto es resultado de considerar al mismo tiempo la escalabilidad y la eficiencia del modelo.
• La razón de implementarlo con NumPy: al implementar el modelo con NumPy, es posible entender de forma más intuitiva su arquitectura y funcionamiento. Esto resulta de gran ayuda para estudiantes e investigadores.
• La incorporación de GQA: GQA es una técnica que ofrece al mismo tiempo ahorro de memoria y mejora del rendimiento, y al aplicarse a todos los modelos de Llama 3 maximiza la eficiencia del modelo.
• La importancia de la caché KV: la caché KV cumple un papel importante en los modelos generativos estilo GPT y permite aumentar considerablemente la eficiencia del cálculo.
• Casos de uso reales: el código de ejemplo permite ejecutar el modelo en la práctica, lo que ofrece una buena oportunidad para comprobar directamente su rendimiento.

[xguru]

Lo publicado en Hacker News estaba en inglés, pero se cambió por un enlace escrito en coreano por el autor original, Likejazz.