Résumé Technique : TetraJet-v2

1. Énoncé du Problème

Le pré-entraînement des Grands Modèles de Langage (LLM) devient de plus en plus prohibitivement coûteux, ce qui motive le besoin d'un entraînement entièrement quantisé en faible précision (FQT). Bien que les formats 4 bits comme NVFP4 offrent des gains d'efficacité significatifs par rapport à FP8 et BF16, parvenir à un entraînement stable et quasi sans perte à cette précision reste un défi critique. Les tentatives précédentes d'entraînement entièrement en 4 bits ont souffert d'une dégradation sévère des performances.

Les auteurs identifient deux goulots d'étranglement d'optimisation principaux qui entravent l'entraînement entièrement en FP4, distincts de l'erreur de quantification simple :

1. Oscillation des Poids : Un mode de défaillance dominant où de petites mises à jour des poids maîtres en haute précision provoquent le basculement répété de leurs représentations quantifiées entre des bins adjacents (par exemple, entre 0 et 0,5). Ce phénomène perturbe l'optimisation, en particulier lorsque les taux d'apprentissage diminuent, empêchant le modèle de converger vers un minimum local stable.
2. Caractéristiques Atypiques (Outliers) : Un petit nombre de canaux d'activation présentent des magnitudes massives, créant des plages dynamiques que le FP4 ne peut pas représenter avec précision. Cela conduit à des facteurs d'échelle élevés qui écrasent la précision des valeurs plus petites au sein du même groupe de quantification.

Les méthodes existantes, telles que la recette NVFP4 de NVIDIA, reposent souvent sur une précision mixte (maintenant certains blocs en BF16) ou échouent à traiter pleinement ces dynamiques d'optimisation dans un contexte de pré-entraînement à partir de zéro.

2. Méthodologie : TetraJet-v2

TetraJet-v2 est une méthode FQT entièrement 4 bits conçue spécifiquement pour NVFP4, appliquant le format aux activations, aux poids et aux gradients dans toutes les couches linéaires. La méthodologie traite les goulots d'étranglement identifiés à travers trois composants principaux :

A. Quantification Double Bloc Non Biaisée pour les Couches Linéaires NVFP4

Pour maximiser la précision dans les contraintes numériques de NVFP4 (qui utilise des groupes de 16 éléments avec une mise à l'échelle 8 bits), les auteurs proposent un schéma de quantification double bloc :

• Structure : Un bloc externe de 128 éléments est mis à l'échelle globalement, tandis que des blocs internes de 16 éléments sont mis à l'échelle localement. Cela évite de calculer la valeur absolue maximale sur l'ensemble du tenseur, améliorant la compatibilité matérielle et la précision par rapport à une mise à l'échelle par tenseur.
• Estimation de Gradient Non Biaisée : La conception de la couche linéaire assure une estimation de gradient non biaisée en :
  • Alignant les tenseurs lors de la rétropropagation (estimant $\hat{X}$ plutôt que $X$).
  • Utilisant un Arrondi Stochastique lors de la rétropropagation pour garantir $E[\text{arrondi}(x)] = x$, empêchant ainsi tout biais dans les attentes de gradient.
  • Adoptant une forme de groupe de poids $1 \times 16$ (non carrée) pour une précision accrue, nécessitant un alignement spécifique pour la matrice de poids transposée.

B. OsciReset : Suppression des Oscillations

Pour traiter l'oscillation des poids, les auteurs introduisent OsciReset, le premier algorithme spécifiquement conçu pour supprimer ce goulot d'étranglement dans le pré-entraînement des LLM.

• Détection : La méthode suit la trajectoire d'optimisation des poids, calculant une métrique de « Risque d'Oscillation » définie comme le rapport entre la distance parcourue par le poids quantifié ($dist_Q$) et la distance parcourue par le poids maître ($dist_M$). Un risque élevé indique qu'un poids bascule de bin malgré de petites mises à jour du poids maître.
• Mécanisme : Lorsque le taux d'apprentissage est faible et que le risque d'oscillation dépasse un seuil, le poids maître de l'élément oscillant est réinitialisé au centre de son bin de quantification actuel.
• Avantage : Contrairement aux méthodes précédentes qui gèlent les poids ou utilisent une régularisation lourde (ce qui nuit à l'optimisation globale), OsciReset permet au poids de continuer à optimiser sur une nouvelle trajectoire moins sujette aux oscillations, sans dégrader immédiatement les performances.

C. OutControl : Contrôle des Atypiques

Pour gérer les valeurs atypiques, TetraJet-v2 emploie une stratégie hybride :

• Transformation Hadamard Aléatoire (RHT) : Appliquée aux calculs de rétropropagation ($dX$ et $dW$) pour distribuer l'énergie des valeurs atypiques plus uniformément sur les canaux, atténuant ainsi la variance de quantification.
• Maintien Statique des Atypiques : Reconnaissant que les canaux atypiques sont structurellement cohérents (certains canaux ont toujours des normes élevées), la méthode identifie les $p\%$ (par exemple 10 %) de canaux atypiques supérieurs lors d'une courte fenêtre d'étalonnage. Ces canaux sont maintenus en précision supérieure (par exemple FP8) pour les passes avant et arrière, tandis que le reste du réseau reste en FP4. Cette sélection statique évite la surcharge computationnelle des ajustements dynamiques par lot.

3. Contributions Clés

1. Quantification Double Bloc Non Biaisée : Une conception novatrice de couche linéaire NVFP4 avec estimation de gradient non biaisée et formes de blocs optimales ($1 \times 16$) pour les LLM, surpassant les recettes NVFP4 existantes.
2. OsciReset : Le premier algorithme efficace pour supprimer l'oscillation des poids dans les LLM durant le pré-entraînement, résolvant une source majeure d'instabilité dans l'entraînement en faible précision.
3. OutControl : Un algorithme de précision mixte combinant RHT pour la rétropropagation avec un maintien statique des atypiques pour les passes avant et arrière afin de préserver la précision.
4. Évaluation Complète : Expériences étendues de pré-entraînement sur des modèles OLMo-2 (jusqu'à 370 millions de paramètres, 212 milliards de tokens) démontrant l'efficacité de la pile complète.

4. Résultats Expérimentaux

Les auteurs ont évalué TetraJet-v2 sur des modèles OLMo-2 (70M, 150M et 370M paramètres) entraînés à partir de zéro.

• Réduction de l'Écart de Performance : TetraJet-v2-full (incorporant toutes les méthodes proposées) a réduit l'écart de performance par rapport à l'entraînement en pleine précision (BF16) d'une moyenne de 51,3 % à travers les tailles de modèles.
• Perplexité (PPL) : Sur le modèle 370M entraîné avec 212 milliards de tokens, TetraJet-v2-full a atteint une PPL de validation de 25,11, contre 23,70 pour BF16 et 26,20 pour la référence NVFP4 de NVIDIA.
• Tâches en Aval : La méthode a obtenu la performance moyenne la plus élevée parmi toutes les méthodes FP4 sur une suite de benchmarks (ARC, BoolQ, MMLU, etc.), un modèle 370M atteignant 43,60 % de précision en moyenne.
• Efficacité : Les implémentations de noyaux CUDA ont atteint une accélération de bout en bout de 1,67× par rapport à FP8 (TransformerEngine) sur du matériel RTX 5090, vérifiant des gains de performance pratiques malgré la complexité ajoutée de la suppression des oscillations et du contrôle des atypiques.
• Études d'Ablation :
  • Le retrait d'OsciReset a conduit à une dégradation sévère des performances et à un rebond de la perte.
  • OutControl a amélioré la précision dans les passes avant et arrière.
  • La méthode s'est révélée robuste face aux choix d'hyperparamètres (par exemple, seuil d'oscillation, moment de début).

5. Signification et Revendications

L'article revendique que TetraJet-v2 représente une étape significative vers la viabilité et la compétitivité du pré-entraînement entièrement en faible précision (4 bits) des LLM par rapport à l'entraînement en haute précision.

• Optimisation vs Représentation : Les auteurs soulignent que l'écart de performance dans l'entraînement en faible précision ne est pas uniquement dû à l'erreur de quantification (représentation) mais est critiquement piloté par des problèmes d'optimisation (oscillations et atypiques) que les méthodes précédentes n'ont pas traités systématiquement.
• Généralisabilité : Contrairement aux techniques antérieures de mitigation des oscillations conçues pour les Transformers de Vision (qui reposent souvent sur un fine-tuning ou un réglage spécifique d'hyperparamètres), OsciReset est présenté comme une méthode simple et efficace applicable au pré-entraînement de LLM à partir de zéro avec un réglage minimal.
• Praticité : En atteignant une accélération de 1,67× par rapport à FP8 tout en réduisant l'écart de précision avec BF16 de plus de 50 %, ce travail démontre que l'entraînement entièrement quantisé peut offrir à la fois une efficacité computationnelle et une qualité de modèle, réduisant potentiellement les coûts matériels et énergétiques de la recherche sur les modèles de base.

Les auteurs reconnaissent les limites concernant l'échelle de leurs expériences (jusqu'à 370 millions de paramètres) et l'absence actuelle de support matériel grand public pour les stratégies de précision mixte comme FP6 × FP4, qu'ils ont explorées comme une direction future potentielle pour des gains de précision supplémentaires.