MXNorm: Reusing MXFP block scales for efficient tensor normalisation

Ce papier présente MXNorm, une méthode de normalisation efficace qui réutilise les échelles de blocs MXFP8 pour estimer la RMS, permettant ainsi une réduction significative des opérations de calcul et des accélérations de vitesse tout en préservant la précision de l'entraînement des modèles Llama 3.

L'essentiel

Le Problème : La Cuisine qui s'emballe

Imaginez que vous dirigez un restaurant de haute technologie (un modèle d'intelligence artificielle comme Llama 3) où les chefs préparent des millions de plats (données) à la seconde.

Depuis quelques années, les cuisiniers (les processeurs) sont devenus incroyablement rapides pour faire les tâches principales : mélanger les ingrédients, couper les légumes et cuire les plats (ce qu'on appelle les multiplications de matrices). C'est comme si on avait remplacé les couteaux en bois par des lasers.

Cependant, il y a un problème : les tâches d'accompagnement, comme peser les ingrédients ou vérifier que la soupe n'est pas trop salée (ce qu'on appelle la normalisation), sont toujours faites avec des balances lentes et des cuillères en bois. Ces tâches ne sont pas devenues plus rapides, alors que le reste de la cuisine a décollé. Résultat : les chefs attendent que la balance fasse son travail, et toute la production ralentit.

La Solution Habituelle : RMSNorm (La Balance Précise)

Pour que les plats aient toujours le même goût, on utilise une technique appelée RMSNorm. C'est comme un chef qui prend chaque assiette, pèse tous les ingrédients un par un, calcule la moyenne, et ajuste le sel. C'est très précis, mais c'est long et cela demande beaucoup d'énergie.

La Nouvelle Idée : MXNorm (L'Estimation Intelligente)

Les chercheurs de Graphcore (l'équipe derrière cet article) ont eu une idée géniale : pourquoi faire deux fois le même travail ?

Dans les cuisines modernes (les puces IA), quand on prépare les ingrédients pour les envoyer dans le four, on doit déjà les "quantifier". C'est-à-dire qu'on les regroupe par paquets de 32 et on leur attribue une étiquette de poids approximative (appelée block scale) pour qu'ils rentrent dans un format plus petit et plus rapide.

MXNorm propose une astuce de maître-chef :

"Au lieu de repeser tout le plat pour ajuster le sel, utilisons simplement l'étiquette de poids qu'on a déjà créée pour le four !"

Au lieu de faire un calcul long et précis pour chaque assiette, MXNorm dit : "On a déjà calculé le poids maximal de chaque paquet d'ingrédients pour les mettre au four. Utilisons ce chiffre pour estimer la moyenne du plat."

L'Analogie de la "Balance de Paquet"

Imaginez que vous devez vérifier si 1000 sacs de pommes sont équilibrés.

1. Méthode ancienne (RMSNorm) : Vous ouvrez chaque sac, comptez chaque pomme, pesez chaque pomme, et faites une moyenne. C'est précis, mais ça prend une heure.
2. Méthode MXNorm : Vous savez déjà que pour les mettre dans un camion, vous avez pesé le sac le plus lourd de chaque groupe de 32. Vous dites : "Bon, si le sac le plus lourd fait 5kg, alors la moyenne du groupe doit être autour de 3kg." Vous faites une estimation très rapide basée sur ce que vous avez déjà fait.

Les Résultats : Est-ce que ça marche ?

Les chercheurs ont testé cette idée sur des modèles géants (Llama 3) :

1. La Précision : Pour les petits modèles, c'est parfait. Pour les très gros modèles, il faut choisir la bonne "recette" d'estimation (comme choisir entre une moyenne simple ou une moyenne pondérée). Avec la bonne recette (p=2), le goût est identique à la méthode ancienne.
2. La Vitesse : C'est là que ça brille. En utilisant cette astuce, le processus de "pesage" devient 2,4 fois plus rapide sur le matériel actuel.
3. L'Économie : Cela réduit la quantité de données à déplacer de 32 fois. C'est comme passer d'un camion de livraison à un vélo pour livrer un petit colis.

En Résumé

MXNorm est une astuce intelligente qui dit : "Ne gaspillons pas notre temps à recalculer ce que nous avons déjà calculé pour une autre étape."

En réutilisant les étiquettes de poids déjà créées pour accélérer la cuisson (la multiplication), on accélère aussi le réglage du goût (la normalisation). C'est une victoire d'efficacité qui permet aux intelligences artificielles de devenir plus rapides et moins gourmandes en énergie, sans perdre en qualité.

C'est un peu comme si, dans votre cuisine, vous utilisiez la même balance pour peser les ingrédients ET pour vérifier la température du four, en une seule passe fluide, au lieu de faire deux voyages séparés.

Résumé technique

1. Problématique et Contexte

L'essor des modèles de langage (LLM) repose sur l'amélioration des capacités des accélérateurs AI, en particulier pour les multiplications matricielles de faible précision (ex: FP8, FP4). Cependant, une disparité croissante existe entre les performances des multiplications matricielles et celles des opérations de réduction (comme le calcul de la norme) et des opérations élémentaires.

• Goulot d'étranglement : Alors que la performance des multiplications matricielles a augmenté de 80x sur les 8 dernières années (grâce aux Tensor Cores), les réductions et les opérations élémentaires n'ont progressé que de 5x à 9x, limitées par la bande passante mémoire et le débit des cœurs CUDA.
• Le cas de la Normalisation : Dans les architectures modernes (Pre-Norm, comme Llama 3), la couche RMSNorm (Root Mean Square Layer Normalization) est exécutée avant la quantification en format MXFP (Microscaling Floating Point).
• Inefficacité actuelle : Le processus standard implique deux passes distinctes :
  1. Calculer la RMS complète sur le tenseur pour normaliser.
  2. Calculer les "block scales" (échelles par bloc) pour la conversion MXFP.
     Cela génère une surcharge inutile en termes de lecture/écriture mémoire et de calculs de réduction, car les deux opérations nécessitent de parcourir les mêmes données le long de la dimension cachée.

2. Méthodologie : MXNorm

Les auteurs proposent MXNorm, un remplacement "plug-and-play" de RMSNorm qui fusionne la normalisation et la conversion MXFP en une seule étape, en réutilisant les statistiques déjà calculées pour la quantification.

Principe Fondamental

L'idée centrale repose sur l'observation que la RMS (moyenne quadratique) d'un tenseur et la moyenne des maxima absolus par bloc (block absmax) sont proportionnelles pour des distributions données (comme une Gaussienne).

• Théorème 1 : Les auteurs démontrent mathématiquement que la moyenne généralisée d'ordre $p$ des maxima absolus des blocs ($G^{(p)}_K$) converge vers la RMS du tenseur multipliée par une constante $c(p, B)$, dépendant uniquement de la taille du bloc $B$ et de l'ordre $p$.
  $$\frac{G^{(p)}_K}{RMS(X)} \xrightarrow{a.s.} c(p, B)$$
• Estimation : Au lieu de calculer la RMS exacte (qui nécessite une réduction sur tout le vecteur), MXNorm estime la RMS en utilisant les maxima absolus des blocs ($m_{tk}$) déjà calculés pour générer les échelles MXFP.
  $$\tilde{\rho}_t \approx c(p, B) \cdot \left( \frac{1}{K} \sum_{k=1}^K m_{tk}^p \right)^{-1/p}$$

Implémentation

1. Fusion : Le calcul de la norme et la mise à l'échelle pour la quantification sont fusionnés. On estime la RMS inversée ($\tilde{\rho}$) à partir des échelles de bloc.
2. Choix de $p$ : Les auteurs testent deux variantes :
   • $p=1$ : Moyenne arithmétique des maxima absolus.
   • $p=2$ : Moyenne quadratique (RMS) des maxima absolus.
3. Gestion des gains (Affine Gain) : Pour intégrer le paramètre d'échelle apprenable $\gamma$ (habituellement appliqué après la normalisation), MXNorm fusionne ce gain avec les poids de la couche linéaire suivante ($W_{fused} = W \cdot \gamma$), évitant ainsi une multiplication élémentaire supplémentaire sur le tenseur quantifié.
4. Rétropropagation : L'algorithme utilise un estimateur "straight-through" réutilisant les gradients de RMSNorm standard, en recachant l'entrée haute précision et l'estimation de l'inverse de la RMS pour le calcul des gradients.

3. Contributions Clés

1. Réduction de la complexité de réduction : MXNorm réduit la taille de la réduction nécessaire pour la normalisation d'un facteur 32x (en fonction de la taille du bloc MXFP, typiquement 32). Au lieu de réduire sur toute la dimension cachée $D$, on réduit sur le nombre de blocs $K = D/B$.
2. Fusion d'opérations : Élimination de la passe séparée pour le calcul de la RMS, permettant de fusionner la normalisation et la conversion MXFP en un seul noyau (kernel).
3. Stabilité théorique et pratique : Démonstration que les bornes de sortie de MXNorm (surtout avec $p=2$) sont comparables à celles de RMSNorm, assurant la stabilité de l'entraînement.
4. Accélération matérielle : Implémentation efficace sur le matériel commercial (Graphcore IPU et NVIDIA GB200) sans modification majeure de l'architecture du modèle.

4. Résultats Expérimentaux

Les auteurs ont validé MXNorm sur l'entraînement préliminaire (pre-training) de modèles Llama 3 de 125M, 1B et 8B paramètres sur le dataset SlimPajama.

Performance et Stabilité

• Modèles petits (125M, 1B) : MXNorm (avec $p=1$ ou $p=2$) montre une sensibilité au taux d'apprentissage et une convergence très similaires à la baseline RMSNorm.
• Modèles grands (8B) :
  • La variante MXNorm ($p=1$) échoue à grande échelle, présentant des pics de perte (loss spikes) dus à une estimation sous-optimale de la RMS en présence de valeurs aberrantes (outliers) et à des bornes de sortie trop larges ($O(K)$).
  • La variante MXNorm ($p=2$) réussit parfaitement. Elle atteint une perte d'entraînement (2.126) quasi identique à la baseline RMSNorm (2.132) et des performances "zero-shot" équivalentes sur les tâches OLMES.
• Conclusion sur la stabilité : La stabilité est corrélée à la capacité de la normalisation à borner les valeurs extrêmes. MXNorm ($p=2$) fournit des bornes ($O(\sqrt{K})$) comparables à RMSNorm ($O(\sqrt{D})$), contrairement à $p=1$.

Accélération (Speedup)

• Kernels isolés : Sur des noyaux fusionnés (Norm + Cast), MXNorm offre un speedup allant jusqu'à 2.4x par rapport à RMSNorm + MXCast, avec une moyenne géométrique de 41.7% pour MXFP8 et 31.2% pour NVFP4.
• Couches complètes (Llama 3 8B) : L'amélioration se traduit par un gain global de 1.3% pour les couches transformer en MXFP8 et 2.6% en NVFP4. Bien que ces chiffres semblent faibles, ils sont significatifs car les opérations de réduction deviennent le goulot d'étranglement dominant lorsque les multiplications matricielles sont extrêmement rapides.

5. Signification et Impact

Ce travail est significatif car il adresse le prochain défi majeur de l'optimisation des LLM : l'efficacité des opérations non-matricielles.

• Optimisation du flux de données : En réutilisant les échelles de blocs déjà calculées pour la quantification, MXNorm élimine une étape coûteuse en bande passante mémoire, ce qui est crucial pour les architectures futures où la mémoire sera le facteur limitant principal.
• Adoption facilitée : La méthode ne nécessite pas de nouveaux hyperparamètres complexes et s'intègre directement dans les pipelines d'entraînement existants (comme TorchTitan).
• Avenir des formats bas de précision : MXNorm démontre que les formats MXFP (Microscaling) peuvent être utilisés non seulement pour les multiplications, mais aussi pour structurer efficacement les opérations de normalisation, ouvrant la voie à une adoption plus large des formats FP4/FP8 dans les modèles de production.

En résumé, MXNorm propose une solution élégante et mathématiquement fondée pour fusionner la normalisation et la quantification, offrant des gains de performance tangibles tout en maintenant la stabilité de l'entraînement de grands modèles de langage.