Llama à partir de zéro : comment implémenter un article sans pleurer

 (blog.briankitano.com)
2 points par GN⁺ 2023-08-10 | 1 commentaires | Partager sur WhatsApp
  • Brian Kitano a construit lui-même une version réduite de Llama avec TinyShakespeare, et conclut que, pour implémenter un article en toute sécurité, il faut partir d’un petit modèle, remplacer les composants un par un, puis entraîner et évaluer à chaque étape.
  • Il commence par mettre en place des fonctions auxiliaires de vérification comme la séparation des données, la génération de batches, l’évaluation de la perte et une fonction de génération, puis vérifie qu’un modèle simple compile et peut s’entraîner avant d’ajouter les composants de Llama.
  • En ajoutant successivement RMSNorm, RoPE et SwiGLU, il vérifie que chaque couche fonctionne comme attendu au moyen des shapes de tenseurs, des propriétés mathématiques et des cartes d’attention.
  • Dans l’attention RoPE, retirer le causal mask fait descendre la perte de validation jusqu’à 0,16, mais dégrade la qualité de génération ; la cause était une fuite d’information due à l’accès aux tokens futurs.
  • Le Llama réduit final comporte 4 blocs et environ 2,37 millions de paramètres, abaisse la perte de validation à environ 1,0, et nécessite aussi de vérifier le flux des gradients ainsi que le planning de taux d’apprentissage.

Commencer petit et construire progressivement sa confiance

  • Pour implémenter un article, l’essentiel est de commencer avec un petit modèle, de remplacer les composants un par un, puis de répéter entraînement et évaluation à chaque changement.
  • Il faut d’abord préparer des fonctions auxiliaires permettant de vérifier le modèle quantitativement :
    • séparation des données
    • boucle d’entraînement
    • visualisation de la perte
    • évaluation de la perte de validation
  • Plutôt que de porter tous les composants de l’article en une seule fois, il prépare aussi une fonction d’évaluation qualitative pour observer les sorties générées par un modèle simple, rapide et déjà familier.
  • Les couches de tenseurs sont vérifiées avec .shape, assert et plt.imshow ; au lieu de commencer par optimiser les multiplications matricielles, il vérifie d’abord les résultats attendus à la main, puis les rend efficaces avec des fonctions torch.
  • Il faut tester en faisant varier la taille de batch, la longueur de séquence et la dimension d’embedding ; un code qui ne fonctionne que pour une seule taille peut casser au moment de l’inférence.

Jeu de données et configuration de base

  • La cible de l’implémentation est une version fortement réduite de Llama de Meta AI, et les données d’entraînement sont TinyShakespeare.
  • Llama est entraîné sur 1,4 T de tokens, mais cette implémentation utilise TinyShakespeare, d’environ 1,11 million de caractères.
  • Le Llama original utilise un tokenizer byte-pair encoding SentencePiece, mais cette implémentation utilise un simple tokenizer au niveau caractère.
    • la taille du vocabulaire est de 65
    • le jeu de données étant petit, le mode de stockage en mémoire n’est pas optimisé séparément
  • Un dictionnaire MASTER_CONFIG gère les réglages du modèle comme vocab_size, batch_size, context_window et d_model.
    • L’objectif est de réduire les constantes et magic numbers, et de rendre le code plus lisible.
  • La fonction get_batches découpe les données en train 80 %, val 10 % et test 10 %, puis génère, depuis un point de départ aléatoire, une entrée x et un label y décalé d’un caractère.

Vérifier la compilation et l’entraînement avec un modèle de base

  • Le premier modèle est SimpleBrokenModel, composé d’un embedding et d’un réseau feed-forward simple :
    • nn.Embedding
    • Linear
    • ReLU
    • Linear
  • Dans l’implémentation d’un article, dire qu’un modèle « fonctionne » signifie qu’il satisfait deux conditions :
    • compilation : les shapes des tenseurs correspondent entre les couches
    • entraînement : la perte baisse effectivement
  • La fonction evaluate_loss échantillonne 10 batches sur les splits train et val, puis calcule la perte moyenne.
  • Après 1000 epochs d’entraînement, SimpleBrokenModel obtenait une perte de validation d’environ 3,94, quasiment inchangée par rapport à la cross-entropy initiale de 4,17.
  • La cause était le passage à F.cross_entropy de valeurs déjà passées par softmax.
    • F.cross_entropy de PyTorch reçoit directement des logits non normalisés.
    • SimpleModel, après retrait du softmax, fait descendre la perte de validation à environ 2,51.
  • Il ajoute ensuite une fonction generate pour inspecter directement les caractères produits par le modèle ; le modèle de base reste imparfait, mais la perte de validation baisse.

Composant Llama 1 : RMSNorm

  • Par rapport au Transformer original, Llama utilise trois modifications architecturales principales :
    • RMSNorm en pre-normalization
    • rotary embeddings
    • fonction d’activation SwiGLU
  • Le Transformer original utilise BatchNormalization, tandis que Llama utilise RMSNorm, qui met le vecteur à l’échelle par la variance sans le centrer.
  • Alors que le Transformer original applique la normalisation à la sortie de la couche d’attention, en post-normalization, Llama utilise une pre-normalization appliquée d’abord à l’entrée.
  • Le RMSNorm implémenté suppose une entrée de shape (batch, seq_len, d_model).
  • Le résultat de RMSNorm est testé grâce à la propriété selon laquelle la norme de couche devient la racine carrée du nombre d’éléments de la couche :
    • assert
    • comparaison ligne par ligne
    • torch.allclose
  • SimpleModel_RMS, qui ajoute RMSNorm au modèle de base, réduit légèrement la perte de validation à environ 2,5015.

Composant Llama 2 : RoPE et causal mask

  • RoPE est une méthode d’encodage positionnel pour Transformers, qui représente la position des tokens par une rotation de l’embedding.
  • get_rotary_matrix génère, pour chaque position, une matrice de rotation à partir de la context window et de la dimension d’embedding.
  • L’implémentation de RoPE est testée avec la propriété suivante :
    • le produit scalaire de deux vecteurs tournés aux positions m et n doit correspondre à une rotation de position relative n-m.
  • RoPEAttentionHead crée w_q, w_k et w_v, applique la rotation RoPE aux query et key, puis utilise F.scaled_dot_product_attention.
  • Il faut faire attention à la différence de shape des tenseurs entre entraînement et inférence :
    • à l’entraînement, on est souvent aligné sur la configuration, par exemple (config['batch_size'], config['context_window'], config['d_model'])
    • à l’inférence, il peut falloir traiter un exemple unique comme (1, 1, config['d_model'])
    • dans forward, l’indexation doit se baser sur la shape obtenue depuis l’entrée, et non sur les valeurs de configuration du modèle
  • Le modèle auquel on ajoute une attention multi-head RoPE sans causal mask voit sa perte de validation chuter brutalement à 0,1623, mais les sorties générées sont mauvaises, du type OOOO... ou IIII....
  • L’inspection de la carte d’attention montre que toutes les positions consultent toutes les positions, provoquant une fuite d’information : le modèle regarde les tokens futurs lors de la prédiction du token suivant.
  • En passant à RoPEMaskedAttentionHead, qui applique is_causal=True à F.scaled_dot_product_attention, l’attention upper triangular correspondant au futur devient presque nulle.
  • Après application du causal mask, la perte de validation est de 2,0815, puis descend à 1,8985 avec un entraînement plus long.

Composant Llama 3 : SwiGLU et empilement de blocs

  • Llama remplace la non-linéarité ReLU par la fonction d’activation SwiGLU.
  • Le SwiGLU implémenté est une Swish-gated linear unit, avec deux transformations linéaires et un paramètre beta apprenable.
  • Le RopeModel qui ajoute SwiGLU dans la partie feed-forward compte 592 706 paramètres, avec une perte de validation d’environ 1,8963.
  • Il crée ensuite LlamaBlock, qui regroupe la configuration suivante dans un seul bloc :
    • RMSNorm en pre-normalization
    • attention multi-head RoPE masquée
    • connexion résiduelle
    • RMSNorm en pre-normalization
    • feed-forward SwiGLU
    • connexion résiduelle
  • Le modèle Llama final est configuré avec n_layers=4 et empile 4 LlamaBlock via un nn.Sequential basé sur OrderedDict.
  • Le modèle final compte 2 370 246 paramètres, et les résultats d’entraînement sont les suivants :
    • après le premier entraînement du modèle à 4 couches, perte de validation de 1,5532
    • après entraînement supplémentaire jusqu’à 10 000 epochs, perte de validation de 1,1479
    • après entraînement additionnel, perte de validation de 0,9997
    • la perte sur un batch du split test est de 1,2358

Résultats de génération et points de débogage

  • Le modèle final produit des noms, retours à la ligne et fragments de mots proches du format de Shakespeare, mais la qualité réelle des phrases reste limitée.
  • La perte cross-entropy peut être interprétée intuitivement du point de vue du choix de tokens :
    • la perte initiale de 4,17 est proche d’un choix aléatoire dans un vocabulaire de taille 65
    • une perte de 1,08 s’interprète comme un choix aléatoire parmi environ 2,9 tokens
  • Le flux des gradients est vérifié avec la fonction show_grads.
    • Elle calcule la proportion de gradients de faible valeur absolue pour chaque paramètre.
    • Si les gradients de la plupart des paramètres ne sont pas proches de 0, le flux est correct.
  • Le Llama original utilise un planning de learning rate Cosine Annealing, mais dans cette implémentation les résultats expérimentaux étaient moins bons.
  • Dans l’expérience Cosine Annealing, même avec une tolerance très faible, l’attention bias recevait très peu de signal ; la raison n’est pas certaine, il est donc plus sûr de commencer simplement dans une implémentation réelle.

À lire aussi

1 commentaires

 
GN⁺ 2023-08-10
Avis de Hacker News

  • Il semble y avoir un bug dans l’implémentation de SwiGLU : dans l’article de référence, le bêta du feed-forward network est une constante, pas une valeur apprenable, et il est défini comme FFnSwiGLU = Swish1...
    D’après l’équation 6 de https://arxiv.org/pdf/2002.05202.pdf
    Dans l’implémentation officielle de llama, le bêta constant est également supprimé : https://github.com/facebookresearch/llama/blob/main/llama/mo...
    Si l’on regarde les lignes "feedforward.1.beta', 0.0" dans les logs du blog, le bêta a dégénéré à 0 pendant l’entraînement, alors qu’il devrait normalement être une constante égale à 1

    • Cela montre à quel point il est difficile d’implémenter correctement un réseau neuronal de type transformeur. On peut se tromper à de nombreuses étapes, et en général cela ne se manifeste que par des performances « légèrement moins bonnes que l’original », ce qui rend le problème difficile à identifier avec certitude
      Il arrive aussi souvent que le réseau s’adapte aux changements, qu’ils soient intentionnels ou non, et qu’après entraînement plusieurs variantes d’architecture se comportent de façon similaire, si bien qu’il peut être ambigu de savoir s’il faut absolument correspondre à l’original
      Une façon de trouver ce type d’erreur consiste à faire correspondre exactement les sorties avec celles d’une implémentation de référence. Comme avec les modèles tiny-random de HuggingFace, même avec des poids aléatoires, la sortie doit être exactement identique ; sinon, c’est le signe d’un bug
      Cette méthode fonctionne toutefois surtout pour les bugs qui apparaissent pendant l’inférence ; les problèmes liés au traitement des données, à l’optimiseur ou qui ne surviennent que pendant l’entraînement sont plus difficiles à détecter
    • Dans les transformeurs, je trouve que les biais ont tendance à mal fonctionner, globalement
      Personnellement, je pense que c’est dû à leur caractère autorégressif et à des propriétés proches des ODE, mais je n’en suis pas certain

  • Le travail est excellent, mais les premiers SimpleBrokenModel et SimpleModel comportent pas mal d’opérations gaspillées. L’ordre est embedding 65 -> 128, linear 128 -> 128, ReLU, linear 128 -> 65, mais il n’y a pas de non-linéarité entre les deux premières couches, et comme elles sont toutes deux linéaires, la deuxième couche linéaire ne sert en fait à rien
    Au final, ce modèle est équivalent à un MLP classique à une seule couche cachée et, en FLOPS, il gaspille 128*128=16k opérations sur un total de 128*128+65*128=24k

    • Je ne suis visiblement pas le seul à être encore en train d’assimiler les non-linéarités. Je me demande si la meilleure correction ici serait d’ajouter ReLU ou SwiGLU entre l’embedding et la première couche linéaire, ou bien de supprimer simplement la couche linéaire
      La couche d’embedding est une structure spéciale qui transforme les indices de tokens en vecteurs d’embedding, donc je ne pense pas qu’on puisse la supprimer

  • Dans l’ensemble, cela illustre bien les principes de base. J’aime particulièrement la phrase « utilisez .shape religieusement. assert et plt.imshow sont vos amis », et il faut toujours vérifier par assert les préconditions et postconditions sur les shapes
    Je me demande aussi si bear ou typeguard prennent en charge ce type de vérification via des décorateurs
    En revanche, la partie « choisissez un modèle petit, simple et rapide, et créez un helper pour l’évaluer qualitativement » veut sans doute dire évaluation quantitative. C’est ce qui permet d’avoir une baseline chiffrée pour comparer avec des techniques plus avancées
    Le conseil d’implémenter les composants de l’article un par un devrait aussi être plus précis. Les articles essaient généralement plusieurs changements en même temps, puis montrent la contribution de chaque élément via des ablations ; il me semble donc préférable de commencer par les changements architecturaux essentiels, puis, dans l’ordre des effets les plus importants observés lors des ablations, d’évaluer chaque changement atomique en respectant les dépendances

    • Au lieu de bear ou typeguard, grâce à https://peps.python.org/pep-0646/, on peut en partie intégrer cela directement dans les annotations de type Python
      Par exemple, on peut représenter la shape par axe dans le type, comme ndarray[float, Dim1, *Shape], et surcharger la shape de retour selon la valeur de axis
    • Je ne connais pas bien PyTorch, mais la dernière fois que j’ai vérifié, ce n’était pas le cas ; Jax prend en charge des vérifications runtime basiques des shapes de matrices via bear / typeguard
      Cela dit, Python aura probablement du mal à être aussi bon que Julia. Le système de types de Julia permet de garantir beaucoup plus facilement que les dimensions des matrices correspondent

  • Je me demande quel est le principe qui pousse à utiliser SwiGLU plutôt que ReLU. Je ne sais pas si les auteurs ont simplement essayé toutes les fonctions non linéaires possibles, ou s’il y a une raison plus profonde

    • Comme souvent en recherche, s’il n’existe pas d’explication claire étayée par une étude rigoureuse, il est probable qu’ils aient simplement fait une recherche par hill climbing au hasard sur des changements d’une ligne qui avaient l’air intéressants, puis se soient arrêtés quand le moment est venu d’écrire l’article et de lancer les ablations

  • Comme bearblog subit une attaque DDoS, je laisse le dépôt : https://github.com/bkitano/llama-from-scratch

  • Du point de vue de quelqu’un qui apprend l’IA, j’ai essayé de résumer simplement les termes cités dans l’article. Un token est un identifiant entier représentant un fragment de texte, et dans les LLM on regroupe des fragments de caractères fréquents dans un vocabulaire de taille limitée.
    La fonction de perte est une valeur qui mesure l’écart entre la prédiction et la réponse attendue ; plus elle est basse, mieux c’est. PyTorch est une bibliothèque pour manipuler des tenseurs et des réseaux de neurones, et un tenseur est un tableau multidimensionnel de nombres incluant scalaires, vecteurs et matrices.
    Un réseau de neurones est une structure de connexions entre neurones avec des poids et des biais, et une couche linéaire est une structure simple où toutes les entrées et sorties sont connectées. ReLU est une fonction d’activation comme Math.max(0, x) ; si l’on empile uniquement des couches linéaires, cela revient finalement à une seule fonction linéaire, donc on ajoute de la non-linéarité pour augmenter la capacité d’apprentissage.
    Un gradient est une quantité de variation numérique calculée pendant l’apprentissage pour rendre le modèle plus précis, et la normalisation par lot est une méthode qui aide l’apprentissage en ajustant les nombres qui circulent. L’encodage positionnel indique sous forme de vecteurs les positions relatives des tokens.
    En Python, l’opérateur @ est un alias de __matmul__ et sert à la multiplication matricielle. Une époque correspond à un apprentissage sur l’ensemble complet du jeu de données, et un batch est le nombre de données injectées en une fois avant la mise à jour des paramètres.
    L’attention est le mécanisme clé qui fait fonctionner les LLM : elle traite les tokens d’entrée en parallèle pour produire des tenseurs intermédiaires, puis les utilise pour générer les tokens de sortie.

    • En dehors du domaine, on peut ne pas savoir ce que signifie « Karpathy ». Présenter Andrej Karpathy avec du contexte, par exemple comme « vulgarisateur scientifique et chercheur », rend plus clair le fait qu’il faut se référer à ses articles ou vidéos.
    • Pour les débutants aussi, il est plus juste de voir un token non pas simplement comme l’identifiant entier d’un fragment de texte, mais plutôt comme un morceau de mot suffisamment fréquent pour être utile en lui-même.
      Par exemple, writ, commun à writing, written et writer, peut devenir un token, et writer peut être tokenisé en writ et er.
      L’embedding est l’étape qui transforme ces tokens en représentations numériques uniques.
    • La composition de fonctions linéaires donne à nouveau une fonction linéaire. Donc si tout est linéaire, empiler plusieurs couches revient à gaspiller toutes les couches sauf une ; pour éviter cela, il faut de la non-linéarité.
    • En plus de la série de vidéos de Karpathy et du repo qui l’accompagne, je me demande s’il existe d’autres ressources ou livres qui ont été particulièrement utiles dans le parcours d’apprentissage.
    • Je me demande ce que fait exactement la normalisation par lot, et en quoi elle aide.

  • S’il existe une implémentation de référence du modèle et des checkpoints, la manière la plus efficace de vérifier que sa propre implémentation est correcte est de charger ces checkpoints et de comparer les sorties.
    Si les sorties ne correspondent pas, c’est généralement qu’un détail de l’implémentation est incorrect, et on peut alors suivre chaque couche de façon systématique pour trouver la différence réelle. On peut même, au passage, découvrir des bizarreries dans l’implémentation existante.
    Cela concerne le modèle lui-même ; l’entraînement est un axe séparé. Cela dit, si les hyperparamètres sont réglés de façon à peu près similaire, quand l’implémentation du modèle est correcte, les choses se passent généralement bien.

  • Les conseils sur la lecture d’articles et le contenu de cet article sont tous deux bons, et je recommande aussi la série Makemore de Karpathy.

  • Les conseils de synthèse sont excellents, et je pense que le conseil d’assert les shapes des tenseurs s’applique à n’importe quelle bibliothèque d’algèbre linéaire généraliste. Quand on écrit du code d’algèbre linéaire complexe, il est très important d’avancer par petites étapes et de coder défensivement.
    Programmer de l’algèbre linéaire dans les langages grand public est pénible parce qu’il n’y a pas de vérification des shapes à la compilation. La shape d’un tenseur devrait faire partie de son type, et essayer de multiplier 3x4 par 3x4 sans transposition devrait tout simplement empêcher la compilation.
    Faire tourner un long calcul pour échouer ensuite sur une opération avec dimensions incompatibles, c’est vraiment le pire.
    Je pense aussi que les tenseurs PyTorch devraient avoir un typage statique de leur périphérique. Aujourd’hui, si l’on essaie de multiplier un tenseur en mémoire CPU par un tenseur en mémoire GPU, on obtient une erreur à l’exécution.