Le modèle Llama 3 implémenté uniquement avec NumPy

• Comprendre précisément l’architecture grâce à une implémentation réellement fonctionnelle du modèle Llama 3

Aperçu

• Le modèle Llama 3 publié par Meta attire l’attention.
• Il affiche une échelle et des performances impressionnantes, avec 24K GPUs, 15T de données d’entraînement, 10M de données d’instructions et 1.3M d’heures GPU.
• L’architecture du modèle n’a pas fondamentalement changé. Llama 3 utilise GQA, mais cela était déjà implémenté dans Llama 2 70B.
• L’implémentation repose uniquement sur NumPy afin de permettre une compréhension intuitive de l’architecture du modèle.
• Le modèle stories15M entraîné par Andrej Karpathy sur l’architecture Llama 2 est converti au format compressé NumPy puis utilisé.

Architecture

• L’architecture du modèle Llama 3 est identique à celle de 42dot LLM.
• Paramètres du modèle :
  • dim: 288
  • n_layers: 6
  • n_heads: 6
  • vocab_size: 32000
  • max_seq_len: 256
  • max_new_tokens: 50

RoPE #1

• Les valeurs cos et sin sont précalculées pour l’embedding RoPE.
• Ces valeurs sont utilisées pour Q et K.
• Le résultat du calcul est multiplié avec np.outer, puis cos et sin sont calculés.

RMSNorm

• RMSNorm normalise les activations par leur Root Mean Square plutôt qu’à partir des statistiques traditionnelles de mini-batch ou de couche.
• Cela fournit une mise à l’échelle cohérente des activations.

QKV

• Pour le calcul de QKV, contrairement à GPT qui applique un matmul unique avant de découper le résultat, Llama dispose de poids distincts pour Q, K et V.
• Chaque valeur est ensuite réorganisée pour le Multi-Head Attention.

RoPE #2

• RoPE possède à la fois des propriétés d’encodage positionnel absolu et relatif.
• Il ne s’applique qu’à Q et K : l’entrée est séparée, multipliée par cos et sin, puis recombinée par addition et soustraction.

Cache KV

• Les modèles génératifs de style GPT utilisent le Masked Attention, ce qui permet un cache KV.
• Comme les résultats précédents restent toujours identiques, K et V sont mis en cache et seul le dernier Q est calculé.

GQA (Grouped-Query Attention)

• GQA est une technique introduite dans Llama 2, qui permet à la fois d’économiser de la mémoire et d’améliorer les performances.
• Dans Llama 3, GQA est appliqué à tous les modèles de 8B et plus.

Scaled Dot-Product Attention

• Chaque attention est calculée via le Multi-Head Attention.
• Le résultat est obtenu par softmax et matmul.

Feed Forward

• Le Feed Forward du modèle Llama utilise 3 couches linéaires, sans biais.
• Il génère une valeur swish, la multiplie par x_V, puis la réduit à nouveau en dimension.

SwiGLU

• SwiGLU améliore les performances du modèle grâce à une combinaison particulière de plusieurs couches feed-forward.

Linear

• La sortie finale calcule uniquement le dernier logit par matmul afin d’accélérer l’exécution.

Génération

• Les tokens sont générés un par un à partir des logits extraits.
• Le processus se divise en phase de prefill et phase de décodage.
• Dans la phase de prefill, toute l’entrée est transmise ; dans la phase de décodage, seul le dernier ID de token est passé pour obtenir le résultat.

Exemple

• Exécution possible comme suit :

  $ python llama3.py "I have a dream"  
  

GitHub

• Le code source complet est disponible sur likejazz/llama3.np.

Références

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

L’avis de GN⁺

• Architecture et performances du modèle Llama 3 : le modèle Llama 3 conserve l’architecture du précédent Llama 2 tout en améliorant nettement les performances. C’est le résultat d’un équilibre entre extensibilité et efficacité.
• Pourquoi une implémentation en NumPy : implémenter le modèle avec NumPy permet de comprendre plus intuitivement sa structure et son fonctionnement. C’est une aide précieuse pour les personnes en apprentissage comme pour les chercheurs.
• L’introduction de GQA : GQA est une technique qui apporte simultanément des gains de mémoire et de performances, et son application à tous les modèles de Llama 3 maximise l’efficacité globale.
• L’importance du cache KV : le cache KV joue un rôle central dans les modèles génératifs de style GPT et permet d’améliorer fortement l’efficacité des calculs.
• Cas d’usage concret : le code d’exemple permet d’exécuter réellement le modèle et constitue une bonne occasion d’en vérifier directement les performances.

[xguru]

Sur Hacker News, le lien publié était en anglais, mais il a été remplacé par celui que l’auteur original, Likejazz, avait rédigé en coréen.