The Llama 3 Herd of Models

Guide de lecture

Problème et motivation

Les systèmes modernes d'intelligence artificielle dépendent fortement de modèles fondationnels massifs, mais les modèles les plus performants (tels que GPT-4) sont restés propriétaires et fermés. Concevoir un modèle de pointe accessible ouvertement impose de surmonter d'immenses goulots d'étranglement d'ingénierie selon trois piliers : la curation des données, la puissance de calcul et la maîtrise de la complexité.

Les auteurs développent Llama 3, une famille de modèles de langage multilingues de 8B, 70B et 405B paramètres prenant en charge des fenêtres de contexte allant jusqu'à 128K tokens, l'utilisation d'outils en zéro-shot et multi-étapes, le raisonnement avancé et la génération de code.

Idées clés

• Maîtrise de la complexité avec une architecture dense : Plutôt que de basculer vers des architectures Mixture-of-Experts (MoE), les auteurs s'en tiennent à un Transformer dense standard (126 couches, dimension de 16 384, 128 têtes d'attention pour le 405B) avec une Grouped-Query Attention (GQA) et RoPE configuré avec une fréquence de base $\theta =500000$ pour maximiser la stabilité du pré-entraînement.
• Mise à l'échelle des données et filtrage rigoureux : Le pré-entraînement passe de 1,8T tokens pour Llama 2 à environ 15,6T tokens issus de domaines diversifiés (~50 % connaissances générales, 25 % mathématiques/raisonnement, 17 % code, 8 % multilingue).
• Lois d'échelle prédictibles : Un cadre de lois d'échelle en deux étapes relie le calcul aux performances sur les tâches en aval (comme ARC Challenge) sur quatre ordres de grandeur avant le lancement de l'entraînement.
• Post-entraînement itératif : Six cycles d'alignement combinant Supervised Fine-Tuning (SFT), Rejection Sampling (RS) et Direct Preference Optimization (DPO), enrichis d'une génération ciblée de données synthétiques.
• Parallélisme 4D et infrastructure : Pré-entraînement du 405B sur un cluster allant jusqu'à 16K GPU H100 à l'aide d'une combinaison de Tensor Parallelism (TP), Context Parallelism (CP), Pipeline Parallelism (PP) et Fully Sharded Data Parallelism (FSDP/DP).

Méthode (avec les équations clés)

1. Formulation de la loi d'échelle

L'allocation de tokens optimale au sens du calcul est modélisée par une loi de puissance dépendant du budget de calcul $C$ :

$$N^{\star }(C)=A{C}^{\alpha }$$

Les ajustements empiriques sur des budgets de calcul compris entre $6\times 10^{18}$ et $10^{22}$ FLOPs ont déterminé $(\alpha ,A)=(0,53,0,29)$. L'extrapolation au budget cible de $3,8\times 10^{25}$ FLOPs a suggéré l'entraînement d'un modèle de 402B paramètres sur 16,55T tokens, guidant ainsi l'architecture finale de 405B paramètres.

2. Extension de la fenêtre de contexte et masquage d'attention

L'extension du contexte de 8K à 128K tokens s'effectue en six étapes de pré-entraînement continu réparties sur environ 800B tokens. Afin d'éviter toute contamination croisée entre les différents documents concaténés au sein d'une même longue séquence de contexte, un masque d'attention de document personnalisé est appliqué.

Le Context Parallelism (CP) découpe les séquences d'entrée en $2\times \text{CP}$ blocs, en exécutant un schéma de communication all-gather sur les représentations des clés ($K$) et des valeurs ($V$). La GQA réduisant la taille de $K$ et $V$ par rapport à celle des requêtes ($Q$), le surcoût de communication de l'all-gather $O(S)$ devient négligeable face au calcul de l'attention $O(S^2)$.

3. Adaptations de la Direct Preference Optimization (DPO)

Lors de l'optimisation des préférences, les tokens de formatage (en-têtes et marqueurs de fin de séquence) sont explicitement masqués de la fonction de perte DPO pour empêcher les répétitions finales ou les arrêts prématurés. De plus, un terme de régularisation auxiliaire basé sur la log-vraisemblance négative (NLL) est ajouté à la séquence choisie :

$${\mathcal{L}}_{\text{reg}}=0,2\cdot {\mathcal{L}}_{\text{NLL}}(\text{chosen})$$

Résultats (avec les chiffres les plus importants)

• Benchmarks du modèle phare 405B :
  • MMLU (5-shot) : 87,3 % (0-shot CoT : 88,6 %)
  • GSM8K (8-shot, CoT) : 96,8 %
  • MATH (0-shot, CoT) : 73,8 %
  • HumanEval (0-shot) : 89,0 %
  • ARC Challenge (0-shot) : 96,9 %
  • IFEval : 88,6 %
• Modèles plus compacts (8B et 70B) : Les deux modèles surpassent leurs équivalents open source comparables (ex. Gemma 2 9B, Mistral 7B, Mixtral 8x22B). Par exemple, Llama 3 8B obtient 73,0 % sur MMLU en 0-shot CoT et 84,5 % sur GSM8K.
• Contexte long : Taux de récupération parfait de 100 % sur Needle-In-A-Haystack simple jusqu'à 128K tokens, et 98,1 % de rappel sur Multi-needle pour le modèle 405B.
• Efficacité du pré-entraînement : La Model FLOPs Utilization (MFU) en BF16 a atteint 38 à 43 % selon les configurations sur 16 384 GPU H100, maintenant un temps d'entraînement effectif supérieur à 90 % sur un instantané de 54 jours malgré 466 interruptions.

Limites et questions ouvertes

• Fiabilité matérielle et fluctuations de charge : À l'échelle de 16K GPU, les interruptions matérielles se sont produites en moyenne plusieurs fois par jour (78 % des interruptions imprévues étant dues à des problèmes matériels/GPU). Les variations simultanées de puissance de plusieurs dizaines de mégawatts lors de la sauvegarde de checkpoints ou des étapes de communication collective mettent à rude épreuve les réseaux électriques des datacenters.
• Écart de robustesse adverse : Bien que résistant sur les benchmarks de paraphrase (PAWS), le modèle montre une dégradation notable sur les benchmarks adverses de raisonnement mathématique et de QA (GSM-Plus, Adv SQuAD) par rapport aux évaluations non adverses.
• Disponibilité multimodale : Les adaptateurs pour la vision et la parole ont démontré des capacités compétitives dans un cadre de recherche, mais sont restés en cours de développement actif et ont été exclus de la première version publique mise à disposition.