Flash Attention est-il stable ? Meta et Harvard ont constaté que les écarts de poids de leur modèle fluctuaient de plusieurs ordres de grandeur.

Meta FAIR s'est associé à Harvard pour fournir un nouveau cadre de recherche permettant d'optimiser les biais de données causés par l'apprentissage automatique à grande échelle.

Comme nous le savons tous, la formation de grands modèles de langage prend souvent des mois et utilise des centaines, voire des milliers de GPU. En prenant comme exemple le modèle LLaMA2 70B, sa formation nécessite un total de 1 720 320 heures GPU. La formation de grands modèles présente des défis systémiques uniques en raison de l’ampleur et de la complexité de ces charges de travail.

Récemment, de nombreuses institutions ont signalé une instabilité au cours du processus de formation lors de la formation des modèles d'IA générative SOTA. Elles apparaissent généralement sous la forme de pics de perte, comme le modèle PaLM de Google, qui se sont produits jusqu'à 20 fois au cours du processus de formation. pointes.

L'écart numérique est à l'origine de cette inexactitude de formation. En raison du coût d'exécution extrêmement élevé de la formation sur de grands modèles de langage, la manière de quantifier l'écart numérique est devenue une question clé.

Dans leurs derniers travaux, des chercheurs de Meta et de l'Université Harvard ont développé une méthode quantitative fondée sur des principes pour comprendre les biais numériques dans l'optimisation de la formation. Ceci est utilisé pour évaluer différentes techniques d'optimisation de pointe et déterminer si elles pourraient introduire des instabilités inattendues lorsqu'elles sont utilisées pour entraîner de grands modèles. Les chercheurs ont découvert que même si les méthodes d’optimisation existantes fonctionnaient bien sur certaines tâches, certains écarts numériques se produisaient lorsqu’elles étaient appliquées à de grands modèles. Ce biais numérique peut créer une instabilité pendant le processus de formation, entraînant une dégradation des performances du modèle. Afin de résoudre ce problème, les chercheurs ont proposé une optimisation basée sur des méthodes quantitatives fondées sur des principes

• Titre de l'article : Is Flash Attention Stable ?
• Lien papier : https://arxiv.org/pdf/2405.02803

Il a été constaté qu'en un seul passage vers l'avant, l'écart numérique de Flash Attention était d'un ordre de grandeur supérieur à l'attention de base de BF16.

Plus précisément, la méthode comprend deux étapes, dont :

• Développer un micro-benchmark pour perturber la précision numérique dans une optimisation donnée
• Évaluer les valeurs numériques grâce à une analyse basée sur les données ; basé sur la distance de Wasserstein Comment le biais se traduit par des changements dans les pondérations du modèle.

Les chercheurs ont analysé la technologie d'optimisation SOTA Flash Attention et ont quantifié l'écart numérique qui peut être introduit. Flash Attention est une technologie largement utilisée pour accélérer les mécanismes d’attention et est souvent considérée comme un goulot d’étranglement du système dans le modèle Transformer. Bien que Flash Attention améliore la vitesse et réduit l'accès à la mémoire, il repose également sur l'optimisation des algorithmes, et l'optimisation des algorithmes peut conduire à une augmentation de l'écart numérique.

Les chercheurs ont émis l'hypothèse que l'ajout de facteurs de redimensionnement pourrait introduire des approximations involontaires, conduisant à des compromis numériques, qui pourraient par la suite affecter la stabilité de l'entraînement.

Ils ont analysé Flash Attention dans le contexte de charges de travail multimodales de texte en image pour déterminer l'importance potentielle des écarts numériques entre Flash Attention et sa ligne de base. En fin de compte, ils introduisent un cadre pour quantifier le biais numérique de l’optimisation de la formation et ses effets en aval.

Les chercheurs ont apporté les deux contributions principales suivantes à la quantification des écarts numériques :

(1) Conception d'un micro-benchmark pour isoler l'impact de la précision numérique sur les écarts numériques.

Le micro-benchmark conçu par les chercheurs est une technique utilisée pour mesurer et quantifier l'écart numérique provoqué par l'optimisation traditionnelle de la boîte noire (telle que Flash Attention). En perturbant des aspects qui ne sont généralement pas disponibles dans les noyaux fournis, ils ont été les premiers à découvrir qu'à faible précision numérique (BF16), Flash Attention a un biais numérique d'environ un ordre de grandeur supérieur à celui de Baseline Attention.

(2) Analyse basée sur les données basée sur la métrique de distance de Wasserstein.

Grâce à cette analyse, les chercheurs contextualisent les écarts numériques observés et forment une limite supérieure de leur impact sur les propriétés du modèle en aval. Dans l'étude de cas des chercheurs, ils ont pu limiter l'impact du biais numérique observé et ont découvert : "Flash Attention a introduit un biais de poids du modèle qui était environ 1/2 à 1/5 fois celui d'un entraînement de faible précision

."

Cette étude souligne l'importance de développer une approche fondée sur des principes pour « non seulement quantifier mais aussi contextualiser l'impact de l'optimisation de la formation sur le biais numérique ». En construisant des proxys pour contextualiser le contexte du biais numérique, dans le but de déduire la probabilité d'effets de modèle en aval (c'est-à-dire , instabilités de formation) qui sont souvent difficiles à mesurer.

Méthodes expérimentales

Les chercheurs ont d'abord développé un micro-benchmark pour isoler et étudier l'écart numérique provoqué par Flash Attention. Comme le montre la figure 2, ils ont réimplémenté numériquement Flash Attention pour analyser différentes précisions numériques et appliquer des mesures d'optimisation potentielles à chaque étape de l'algorithme.

Figure 2 : Résumé de la conception du microbenchmark.

Cela est nécessaire car le noyau Flash Attention ne prend actuellement en charge que les formats numériques FP16 et BF16. Ce noyau est également un appel API wrapper pour le code CUDA, ce qui rend difficile la perturbation de l'algorithme pour examiner l'impact du biais numérique.

En revanche, leur conception de micro-référence permet une saisie et une modification précises au sein de l'algorithme. Les chercheurs ont vérifié le microbenchmark par rapport au noyau Flash Attention d'origine.

Ils ont en outre conçu une technique pour comparer le résultat de la matrice Attention à chaque étape de l'exécution du modèle. Et modifié le code du modèle pour calculer l'attention de base et l'attention flash à chaque fois que l'attention est appelée, ce qui permet une comparaison précise de la matrice de sortie pour la même matrice d'entrée.

Pour mettre cela en contexte, nous avons également utilisé les métriques de différence maximale et de distance de Wasserstein pour quantifier la différence de poids des modèles tout au long de l'entraînement en utilisant des séries d'entraînement identiques et indépendantes.

Pour l'expérience de formation, les chercheurs ont utilisé une charge de travail d'IA générative (c'est-à-dire un modèle texte en image) qui convertit la saisie de texte en images. Ils ont recyclé le modèle à l'aide de l'ensemble de données Shutterstock et ont mené l'expérience sur un cluster de GPU NVIDIA 80 Go A100.

Quantification des écarts numériques grâce à des micro-benchmarks

Les chercheurs ont d'abord analysé l'impact de Flash Attention dans le processus de passe avant. Ils ont utilisé des micro-benchmarks pour examiner l’impact de différentes précisions numériques sur la matrice de sortie calculée par Attention, à condition que les vecteurs de requête, de clé et de valeur initialisés de manière aléatoire soient les mêmes.

Comme le montre la figure 3, lorsque les chercheurs utilisent différents formats numériques allant de BF16 à FP64, l'écart numérique entre Flash Attention et Baseline Attention diminue à mesure que le nombre de chiffres de la mantisse augmente. Cela suggère que la différence numérique est due à l'approximation inhérente au fait d'avoir moins de chiffres de mantisse.

Figure 3 : L'effet du format numérique sur la déviation numérique de Flash Attention.

Ensuite, le chercheur a défini une « valeur d'or » pour l'attention de base dans le format numérique FP64 pour une comparaison standard, puis a comparé la sortie Attention dans différents formats numériques avec cette valeur (comme le montre la figure 4).

Figure 4 : Comparaison de la « valeur or » de l'attention de base sous le FP64.

Les résultats montrent que l'écart numérique de Flash Attention est environ 10 fois supérieur à celui de Baseline sous BF16.

Pour analyser plus en détail cet écart numérique observé, les chercheurs ont scanné la longueur de séquence de la matrice tout en gardant constantes la taille des tuiles et la taille de la SRAM (comme le montre la figure 5).

Figure 5 : L'effet de la longueur de la séquence sur l'écart numérique de Flash Attention.

Comme le montre la figure, à mesure que la longueur de la séquence augmente, qu'elle soit mesurée par (a) la limite supérieure de la différence maximale, ou par (b) la moyenne et l'écart type de la différence, la différence entre Flash Attention et Baseline Attention Les écarts numériques augmentent.

De plus, les chercheurs utilisent également des conceptions de micro-référence pour mener des expériences avec différentes optimisations afin de mieux comprendre l'impact des écarts numériques (comme le montre la figure 6).

La figure 6a montre comment l'échange de l'ordre des dimensions des blocs entraîne une augmentation de la différence numérique entre Flash Attention et Baseline Attention. D'autres perturbations sur la figure 6b, telles que la limitation de la taille des carreaux aux carrés, n'ont aucun effet sur le biais numérique. La figure 6c montre que plus la taille du bloc/tuile est grande, plus l'écart numérique est petit.

Figure 6 : Changements d'algorithme et leur impact sur les écarts numériques observés.

Comprendre le biais numérique dû aux différences de poids

Bien que Flash Attention puisse provoquer un biais numérique dans la sortie d'attention lors de la passe avant, le but ultime de cette étude est de déterminer si cela se produira pendant l'entraînement du modèle et produira des effets pour déterminer si cela contribue à l’instabilité de la formation.

Par conséquent, les chercheurs espèrent quantifier si Flash Attention modifie le modèle pendant l'entraînement, c'est-à-dire si la différence de sortie d'attention observée ci-dessus se reflète dans les poids du modèle mis à jour pendant l'entraînement.

Les chercheurs ont utilisé deux indicateurs pour mesurer la différence de poids des modèles entre les modèles formés à l'aide de Baseline Attention et les modèles formés à l'aide de Flash Attention. La différence maximale est d'abord calculée, c'est-à-dire en trouvant la valeur absolue de la différence entre les matrices de poids et en prenant la valeur maximale, obtenant ainsi la limite supérieure de l'écart, comme suit :

Alors que la différence maximale fournit une limite supérieure de l'écart numérique, mais elle ne prend pas en compte la distribution de chaque matrice. Par conséquent, les chercheurs quantifient les différences de poids grâce à la distance de Wasserstein, qui est une mesure courante de similarité entre les tenseurs. Bien que légèrement plus complexe sur le plan informatique, Wasserstein Distance inclut des informations sur la forme de la distribution tensorielle pour mesurer la similarité. La formule de calcul se résume ainsi :

Plus la valeur est faible, plus la similarité entre les matrices est élevée.

À l'aide de ces deux mesures, les chercheurs ont ensuite quantifié l'évolution des poids du modèle de Flash Attention par rapport à l'attention de base tout au long du processus de formation :

Selon la distance de Wasserstein et la différence maximale, pour deux indicateurs, pendant Tout au long du processus de formation, l'ajout de Flash Attention modifie le poids du modèle, et à mesure que la formation se poursuit, cette différence ne fera que devenir de plus en plus grande. Cela montre que le modèle formé à l'aide de Flash Attention est différent du modèle formé à l'aide de Baseline Attention. Le même modèle formé a convergé vers un modèle différent.

Cependant, la formation est un processus stochastique et certains changements dans la structure du modèle peuvent produire des résultats similaires en termes d'effets en aval et de précision. Ceci est remarquable même si les poids des modèles formés avec Flash Attention et Baseline Attention sont différents.

Former complètement un modèle et évaluer sa précision est une tâche coûteuse et gourmande en ressources, en particulier pour les grands modèles dont la formation prend des mois.

Le chercheur a configuré un proxy pour explorer :

(a) Quelle est l'importance de ces changements de poids ?

(b) Cela peut-il être lié aux changements de poids standard dans d'autres optimisations d'entraînement largement adoptées ?

Afin d'atteindre cet objectif, les chercheurs ont conçu une série d'expériences pour comparer l'évolution de la différence de poids au cours du processus d'entraînement dans différents scénarios.

En plus de comparer le processus d'entraînement à l'aide de Flash Attention et de Baseline Attention, ils ont également quantifié la différence de poids au cours du même processus d'entraînement où les poids étaient initialisés à différentes valeurs aléatoires au début de l'entraînement. Cela fournit une limite, car l'initialisation de poids aléatoire est une technique courante et produit souvent des résultats équivalents.

De plus, les chercheurs ont également mesuré les changements dans les poids des modèles entraînés avec différentes précisions. La précision numérique (c'est-à-dire FP16 par rapport à FP32) a le potentiel de provoquer des changements en aval, ce qui sert de limite supérieure à l'importance des pondérations Flash Attention.

Comme le montre la figure 8, on peut constater que le taux de changement de biais de poids du modèle à l'aide de Flash Attention est comparable ou inférieur au taux de changement de biais de poids des différentes initialisations du modèle (notez la pente des courbes rouge et bleue) .

De plus, le taux de changement de poids lors de l'utilisation du FP16 et du FP32 est plus élevé et le changement est plus important que lorsque différents modèles sont initialisés.

Ces résultats fournissent une approximation et montrent : "Bien que Flash Attention présente un biais numérique, il sera limité par l'initialisation aléatoire du modèle et l'entraînement de faible précision. Et le biais de poids du modèle introduit est d'environ 10 % lors d'un entraînement avec une faible précision. . 1/2 à 1/5 fois "

Figure 8 : Différence de poids relative pendant l'entraînement mesurée à l'aide de la métrique de distance Wasserstein.

Pour plus de détails sur la recherche, veuillez vous référer à l'article original.

Ce qui précède est le contenu détaillé de. pour plus d'informations, suivez d'autres articles connexes sur le site Web de PHP en chinois!