Improved Scaling for Fast Mode of Ozaki Scheme II

Résumé Technique : Amélioration de la mise à l'échelle du mode Rapide du Schéma d'Ozaki II

Énoncé du Problème
Le Schéma d'Ozaki II est une méthode pour simuler des multiplications de matrices de haute précision (ex: FP64/FP32) en utilisant du matériel de bas niveau de précision entier (ex: Tensor Cores INT8) basé sur le Théorème des Restes Chinois (CRT). Le schéma nécessite de mettre à l'échelle les matrices d'entrée de haute précision en matrices entières de telle sorte que le produit de ces entiers, lorsqu'ils sont reconstitués via le CRT, ne dépasse pas un seuil spécifique ($P/2$). Si ce seuil est dépassé, la récupération par CRT échoue, entraînant une perte de précision catastrophique.

Le schéma propose deux modes pour déterminer les facteurs d'échelle ($\mu$ et $\nu$) :

1. Mode Précis (Accurate Mode) : Utilise une multiplication de matrices (GEMM) entière INT8 auxiliaire pour estimer les facteurs d'échelle, garantissant une grande précision mais entraînant un surcoût computationnel important.
2. Mode Rapide (Fast Mode) : Utilise l'inégalité de Cauchy-Schwarz pour estimer les facteurs d'échelle sans la GEMM auxiliaire, offrant un débit plus élevé mais historiquement une précision moindre.

Cet article identifie une faille critique dans la formule de mise à l'échelle existante du Mode Rapide : elle manque d'invariance d'échelle. Lorsque les matrices d'entrée sont multipliées par une constante scalaire $\alpha$, la largeur de bits effective des matrices entières mises à l'échelle résultantes change de manière imprévisible.

• Pour un $\alpha$ grand, la largeur de bits effective diminue, provoquant une dégradation de la précision.
• Pour un $\alpha$ petit, le produit entier mis à l'échelle peut dépasser le seuil du CRT ($P/2$), provoquant un échec complet de la récupération.
  Des expériences numériques sur des matrices remplies de uns et des matrices aléatoires ont démontré que, pour certaines plages scalaires, le Mode Rapide ne parvient pas à récupérer le résultat correct, alors que le Mode Précis reste stable.

Méthodologie
Les auteurs proposent une formule de mise à l'échelle révisée pour le Mode Rapide, dérivée directement de la condition d'unicité du CRT via l'inégalité de Cauchy-Schwarz.

1. Dérivation Théorique :
   Les auteurs reformulent la condition pour les facteurs d'échelle $\mu_i$ et $\nu_j$. Au lieu de l'heuristique précédente, ils définissent une condition suffisante où :
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   En résolvant pour $\mu_i$ et $\nu_j$ à l'aide de logarithmes en base 2 et en tenant compte des bornes d'erreur absolue de la fonction log2f fournie par les API mathématiques des GPU, ils dérivent de nouvelles formules :
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   où $P'_{prop}$ est une constante précalculée et $\widehat{\log_2}$ est une approximation bornée du logarithme.

2. Preuve d'Invariance d'Échelle :
   Les auteurs démontrent analytiquement que la nouvelle formule est invariante d'échelle. Si la matrice d'entrée $A$ est mise à l'échelle par $\alpha = 2^s$, le terme $\log_2(\alpha)$ s'annule dans le calcul de la largeur de bits effective de la matrice entière mise à l'échelle. Par conséquent, la largeur de bits reste constante quel que soit l'ordre de grandeur de l'entrée, empêchant la dégradation de la précision et l'échec du CRT observés dans l'original Mode Rapide.

3. Implémentation :
   La méthode proposée (OS II-prop) remplace la formule de mise à l'échelle dans le Mode Rapide de la bibliothèque GEMMul8. Elle conserve la complexité computationnelle $O(mk)$ et $O(kn)$ du Mode Rapide original, n'introduisant aucun surcoût supplémentaire par rapport au Mode Rapide original et restant plus rapide que le Mode Précis (qui nécessite une GEMM supplémentaire).

Principales Contributions

• Identification de la Limite : L'article démontre que la formule de mise à l'échelle du Mode Rapide existante manque d'invariance d'échelle, entraînant une dégradation de la précision et un échec de la récupération par CRT pour des mises à l'échelle d'entrée spécifiques.
• Formule Révisée : Une nouvelle formule de mise à l'échelle est proposée, qui garantit l'invariance d'échelle et assure que la condition d'unicité du CRT est toujours satisfaite.
• Zéro Surcoût : La nouvelle formule n'introduit aucun coût de calcul supplémentaire par rapport au Mode Rapide original.
• Évaluation Complète : La méthode est évaluée sur un GPU NVIDIA GH200 pour les multiplications de matrices en double précision (DGEMM) et en simple précision (SGEMM).

Résultats
Les expériences ont été menées à l'aide de matrices aléatoires avec des écarts de magnitude des éléments ($\phi$) variables et des dimensions de matrice ($k \in \{1024, 16384\}$).

• Précision : La méthode proposée (OS II-prop) atteint une précision comparable au Mode Précis (OS II-accu) pour toutes les valeurs de $\phi$. En revanche, le Mode Rapide original (OS II-fast) a montré une dégradation significative de la précision et des échecs pour les valeurs de $\phi$ élevées (cas mal conditionnés).
• Débit : OS II-prop maintient un débit comparable au Mode Rapide original. Bien qu'OS II-prop soit légèrement plus lent que le Mode Rapide original dans certains scénarios avec peu de modules (en raison de moins d'éléments nuls dans les entiers mis à l'échelle), il surpasse significativement le Mode Précis en termes de débit.
• Compromis : Dans l'analyse du compromis précision-débit, OS II-prop domine OS II-accu (débit plus élevé pour la même précision) et surmonte les limitations de précision de OS II-fast. C'est la seule méthode capable de surpasser simultanément cuBLAS en précision et en débit pour une large gamme de conditions.

Signification
L'article affirme que la méthode proposée résout la limitation fondamentale de précision du Mode Rapide du Schéma d'Ozaki II sans sacrifier son avantage de performance. En garantissant l'invariance d'échelle, la méthode garantit que la récupération par CRT réussit quel que soit l'étalonnage de l'entrée, offrant un meilleur équilibre entre précision et performance. Cela permet aux utilisateurs d'obtenir des résultats de haute précision sur du matériel de bas niveau de précision avec la vitesse du Mode Rapide et la fiabilité du Mode Précis, surmontant ainsi les contraintes de débit du Mode Précis et les contraintes de précision de l'original Mode Rapide.

Les auteurs notent que bien que le présent travail se concentre sur les moteurs de matrice INT8 pour DGEMM et SGEMM, l'approche est applicable aux moteurs FP8 et aux multiplications de matrices complexes (ZGEMM/CGEMM), ce qui est laissé pour des travaux futurs.