Compter le calcul

Premier article de la série Smallem : les lois d'échelle sur une seule carte, qui prolonge Smallem : un LLM en Rust et CUDA, à partir de zéro.

La première série a toujours entraîné le même modèle : 29 millions de paramètres, huit couches, des vecteurs de 512. Elle a montré que la qualité se paie en calcul. Mais elle n'a jamais posé la question qui vient ensuite : pour un budget de calcul donné, quel modèle faut-il entraîner ? Un modèle deux fois plus petit, entraîné deux fois plus longtemps, aurait-il fait mieux ?

Les lois d'échelle répondent à cette question par la mesure. On entraîne beaucoup de modèles de tailles différentes, avec des budgets différents, et l'on cherche les régularités : la perte suit-elle une loi simple de la taille et des données ? Si oui, on peut prédire, avant de lancer un entraînement, ce qu'il donnera.

Cette série refait ce travail à notre échelle, sur une seule RTX 3090, en une vingtaine d'heures de calcul. La première chose à faire est de savoir mesurer ce calcul.

Les lettres

Les comptes qui suivent utilisent ces lettres :

L'unité : l'opération

Le budget d'un entraînement ne se mesure pas en pas, ni en heures : un pas d'un petit modèle ne coûte pas un pas d'un grand, et une heure dépend du matériel. Il se mesure en opérations, en FLOP.

Presque tout le calcul d'un Transformer est dans ses produits matriciels. Un produit [m, k] · [k, n] fait m × k × n multiplications-additions, soit 2 m k n opérations. Les normalisations, les activations et le softmax font quelques opérations par élément : à côté, c'est négligeable.

Pour un token qui traverse une couche du modèle moderne de largeur C, avec un feed-forward SwiGLU de largeur h ≈ 8C/3 :

projections de l'attention   4 × C²          requêtes, clés, valeurs, sortie
scores et mélange            2 × T × C       chaque requête face aux T positions
feed-forward                 3 × C × h       porte, montée, descente

et à la fin, C × V pour la sortie, face au vocabulaire. C'est une passe avant. La passe arrière fait, pour chaque produit, deux produits : le gradient par rapport à chacun des facteurs. Un pas d'entraînement coûte donc trois passes avant.

Le module flops.rs fait ce compte, couche par couche, pour toutes les architectures de la première série sauf deux : le mélange d'experts et la prédiction multi-token.

Vérifier le compte

Un compte écrit à la main peut oublier un produit. Pour le vérifier, un backend de test enveloppe le backend CPU et compte les multiplications-additions de chaque produit matriciel qu'on lui demande. Un vrai pas d'entraînement y passe, et le test compare.

Il y avait un écart, et il était juste : la passe arrière de l'attention recalcule les poids de l'attention au lieu de les avoir gardés. C'est un produit de plus par couche, que la définition du modèle ne demande pas. L'article 17 de la première série fait la même chose sur la carte, et pour la même raison : recalculer coûte moins cher que garder en mémoire. Une fois ce recalcul compté à part, le compte tombe juste, au produit près, pour cinq architectures : GPT-2, moderne, clés partagées, couches traversées deux fois, attention latente.

On compte donc ici le calcul du modèle, celui que sa définition demande. Ce que font nos kernels s'en écarte un peu : l'attention fusionnée saute les tuiles entièrement dans le futur, ce qui réduit le nombre d'opérations, et recalcule les scores dans sa passe arrière, ce qui augmente le nombre d'opérations.

Ce qu'oublie 6ND

L'approximation habituelle du budget est F ≈ 6 × N × D : N paramètres, D tokens, deux opérations par paramètre et par token à l'aller, quatre au retour (Kaplan et al., 2020).

Elle oublie l'attention : les scores et le mélange ne sont pas des paramètres, mais coûtent 2 × T × C multiplications-additions par token et par couche. Pour notre modèle de 29 millions de paramètres, avec des séquences de 512 tokens, cela fait 14 % de plus que 6ND. L'écart croît avec le contexte, et décroît avec la largeur.

Et quels paramètres compter dans N ? La table d'embedding de l'entrée n'est qu'une lecture : un token y choisit sa ligne, sans aucune multiplication. Mais notre modèle partage cette table avec sa sortie, qui, elle, fait C × V multiplications par token. Les deux articles de référence ne font pas le même choix. Kaplan et al. comptent les paramètres hors embedding, et trouvent ainsi des lois plus nettes. Hoffmann et al. (2022, Chinchilla) comptent tous les paramètres. Le choix change les exposants, et il pèse d'autant plus que le modèle est petit. Notre plus petit modèle a 1,45 million de paramètres, dont 1,05 million dans sa table d'embedding. Le plus grand a 91 millions, dont 6,3. Nous ajusterons les deux.

Le rendement de la carte

Combien de ces opérations la carte fait-elle par seconde, et quelle part de son maximum ? C'est le MFU, le rendement en opérations du modèle (Model FLOPs Utilization, Chowdhery et al., 2022, PaLM). On n'y compte que les opérations que le modèle exige, pas celles que la carte refait en plus quand elle recalcule des activations au lieu de les garder.

La fiche de NVIDIA donne à la RTX 3090 71 téraopérations par seconde sur ses Tensor Cores en bf16, avec accumulation en f32 : c'est le calcul de nos produits en précision mixte. Un programme mesure le débit de chaque configuration en entraînement, lots de 16 séquences de 512 tokens, Muon, précision mixte :

Le plan estimait le rendement du modèle xl à 9 téraopérations par seconde, par 6ND : le compte exact donne 10,0, soit 14 % du maximum de la carte. Les petits modèles font pire : leurs produits sont trop petits pour occuper la carte, et le reste du pas, normalisations, optimiseur, lancements, pèse plus lourd. Et l'écart avec 6ND est plus grand pour eux : l'attention, qui ne dépend que de la largeur et du contexte, y pèse plus, relativement aux paramètres.

Ces 14 % sont loin des 40 à 50 % que visent les grands entraînements. La première série a déjà dit pourquoi : nos activations restent en f32, et beaucoup de nos kernels sont limités par la mémoire. Pour ce plan, cela ne change rien : le budget se compte en opérations du modèle, et la carte ne fait que fixer combien d'heures il coûte.

Ce qu'il faut retenir

• Le budget d'un entraînement se compte en opérations : trois passes avant, deux opérations par multiplication-addition.
• Le compte se vérifie contre un vrai pas : il y manquait le recalcul de l'attention, que la définition du modèle ne demande pas.
• 6ND oublie l'attention, 14 % à 512 positions, et laisse ouvert le choix des paramètres à compter.

Article suivant : Une famille de modèles.