Improved Scaling for Fast Mode of Ozaki Scheme II

Technische Samenvatting: Verbeterde Schaling voor de Fast Mode van het Ozaki Schema II

Probleemstelling
Ozaki Schema II is een methode om hoogprecisie matrixvermenigvuldiging (bijv. FP64/FP32) te emuleren met behulp van laagprecisie integer-hardware (bijv. INT8 Tensor Cores) gebaseerd op het Chinese Reststelling (CRT). Het schema vereist het schalen van hoogprecisie inputmatrices naar integer-matrices zodanig dat het product van deze integers, wanneer gereconstrueerd via CRT, een specifieke drempel ($P/2$) niet overschrijdt. Als deze drempel wordt overschreden, faalt de CRT-reconstructie, wat leidt tot catastrofaal verlies van nauwkeurigheid.

Het schema biedt twee modi voor het bepalen van schalingsfactoren ($\mu$ en $\nu$):

1. Accurate Mode: Gebruikt een hulp-INT8 matrixvermenigvuldiging (GEMM) om schalingsfactoren te schatten, wat zorgt voor een hoge nauwkeurigheid maar aanzienlijke computationele overhead met zich meebrengt.
2. Fast Mode: Gebruikt de Cauchy-Schwarz ongelijkheid om schalingsfactoren te schatten zonder de hulp-GEMM, wat een hogere doorvoer biedt maar historisch gezien een lagere nauwkeurigheid heeft.

Dit artikel identificeert een kritiek gebrek in de bestaande Fast Mode schalingsformule: deze mist schaalinvariantie. Wanneer inputmatrices worden vermenigvuldigd met een scalaire constante $\alpha$, verandert de effectieve bitbreedte van de resulterende geschaalde integer-matrices op onvoorspelbare wijze.

• Voor een grote $\alpha$ neemt de effectieve bitbreedte af, wat leidt tot degradatie van de nauwkeurigheid.
• Voor een kleine $\alpha$ kan het geschaalde integer-product de CRT-drempel ($P/2$) overschrijden, wat leidt tot volledige mislukking van de reconstructie.
  Numerieke experimenten op matrices met enkel enen en op willekeurige matrices toonden aan dat de Fast Mode voor bepaalde scalaire bereiken niet in staat is het juiste resultaat te herstellen, terwijl de Accurate Mode stabiel blijft.

Methodologie
De auteurs stellen een herziene schalingsformule voor de Fast Mode voor die rechtstreeks is afgeleid van de CRT-uniciteitsvoorwaarde via de Cauchy-Schwarz ongelijkheid.

1. Theoretische Afleiding:
   De auteurs herformuleren de voorwaarde voor de schalingsfactoren $\mu_i$ en $\nu_j$. In plaats van de vorige heuristiek, definiëren zij een voldoende voorwaarde waarbij:
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   Door $\mu_i$ en $\nu_j$ op te lossen met behulp van basis-2 logaritmen en rekening te houden met de absolute foutmarges van de log2f functie geleverd door GPU math API's, leiden zij nieuwe formules af:
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   waarbij $P'_{prop}$ een vooraf berekende constante is en $\widehat{\log_2}$ een begrensde benadering van de logaritme is.

2. Bewijs van Schaalinvariantie:
   De auteurs demonstreren analytisch dat de nieuwe formule schaalinvariant is. Als de inputmatrix $A$ wordt geschaald met $\alpha = 2^s$, dan valt de term $\log_2(\alpha)$ weg in de berekening van de effectieve bitbreedte van de geschaalde integer-matrix. Bijgevolg blijft de bitbreedte constant, ongeacht de grootte van de input, wat de degradatie van nauwkeurigheid en de CRT-fout die in de oorspronkelijke Fast Mode werden waargenomen, voorkomt.

3. Implementatie:
   De voorgestelde methode (OS II-prop) vervangt de schalingsformule in de Fast Mode van de GEMMul8 bibliotheek. Het behoudt de $O(mk)$ en $O(kn)$ computationele complexiteit van de oorspronkelijke Fast Mode, introduceert geen extra overhead vergeleken met de oorspronkelijke Fast Mode en blijft sneller dan de Accurate Mode (die een extra GEMM vereist).

Belangrijkste Bijdragen

• Identificatie van Beperking: Het artikel toont aan dat de bestaande Fast Mode schalingsformule een gebrek aan schaalinvariantie vertoont, wat leidt tot degradatie van de nauwkeurigheid en CRT-reconstructiefouten voor specifieke input-schalingen.
• Herziene Formule: Een nieuwe schalingsformule wordt voorgesteld die schaalinvariantie garandeert en ervoor zorgt dat de CRT-uniciteitsvoorwaarde altijd wordt voldaan.
• Nul Overhead: De nieuwe formule introduceert geen extra computationele kosten ten opzichte van de oorspronkelijke Fast Mode.
• Uitgebreide Evaluatie: De methode wordt geëvalueerd op een NVIDIA GH200 GPU voor zowel Double-Precision (DGEMM) als Single-Precision (SGEMM) matrixvermenigvuldigingen.

Resultaten
Experimenten werden uitgevoerd met willekeurige matrices met variërende element-grootte spreidingen ($\phi$) en matrixdimensies ($k \in \{1024, 16384\}$).

• Nauwkeurigheid: De voorgestelde methode (OS II-prop) bereikt een nauwkeurigheid die vergelijkbaar is met de Accurate Mode (OS II-accu) voor alle waarden van $\phi$. In tegenstelling hiertoe vertoonde de oorspronkelijke Fast Mode (OS II-fast) significante degradatie van de nauwkeurigheid en falen voor grote $\phi$-waarden (ill-conditioned gevallen).
• Doorvoer: OS II-prop behoudt een doorvoer die vergelijkbaar is met de oorspronkelijke Fast Mode. Hoewel OS II-prop in specifieke scenario's met een klein aantal moduli iets langzamer is dan de oorspronkelijke Fast Mode (door minder nul-elementen in de geschaalde integers), is het aanzienlijk sneller dan de Accurate Mode qua doorvoer.
• Afweging: In de analyse van de nauwkeurigheid-doorvoer afweging domineert OS II-prop de OS II-accu (hogere doorvoer voor dezelfde nauwkeurigheid) en overwint het de nauwkeurigheidsbeperkingen van de OS II-fast. Het is de enige methode die in staat is om simultaan cuBLAS te overtreffen in zowel nauwkeurigheid als doorvoer voor een breed scala aan condities.

Betekenis
Het artikel beweert dat de voorgestelde methode de fundamentele nauwkeurigheidsbeperking van de Fast Mode in Ozaki Schema II oplost zonder het prestatievoordeel op te offeren. Door schaalinvariantie te garanderen, verzekert de methode dat de CRT-reconstructie slaagt, ongeacht de input-schaling, wat een superieur evenwicht biedt tussen nauwkeurigheid en prestaties. Dit stelt gebruikers in staat om hoogprecisie resultaten te verkrijgen op laagprecisie hardware met de snelheid van de Fast Mode en de betrouwbaarheid van de Accurate Mode, waardoor de doorvoerbeperkingen van de Accurate Mode en de nauwkeurigheidsbeperkingen van de oorspronkelijke Fast Mode effectief worden overwonnen.

De auteurs merken op dat hoewel het huidige werk zich richt op INT8 matrix engines voor DGEMM en SGEMM, de aanpak toepasbaar is op FP8 engines en complexe matrixvermenigvuldigingen (ZGEMM/CGEMM), welke zijn voorbehouden aan toekomstig werk.