Improved Scaling for Fast Mode of Ozaki Scheme II

Technische Zusammenfassung: Verbesserte Skalierung für den Fast-Modus des Ozaki-Schemas II

Problemstellung
Das Ozaki-Schema II ist eine Methode zur Emulation hochpräziser Matrixmultiplikationen (z. B. FP64/FP32) unter Verwendung von Low-Precision-Integer-Hardware (z. B. INT8-Tensor-Cores) basierend auf dem Chinesischen Restsatz (Chinese Remainder Theorem, CRT). Das Schema erfordert die Skalierung hochpräziser Eingangsmatrizen in Integer-Matrizen, sodass das Produkt dieser Integer, wenn sie via CRT rekonstruiert werden, einen spezifischen Schwellenwert ($P/2$) nicht überschreitet. Wenn dieser Schwellenwert überschritten wird, schlägt die CRT-Rekonstruktion fehl, was zu einem katastrophalen Genauigkeitsverlust führt.

Das Schema bietet zwei Modi zur Bestimmung der Skalierungsfaktoren ($\mu$ und $\nu$):

1. Accurate Mode (Genauer Modus): Verwendet eine zusätzliche INT8-Matrixmultiplikation (GEMM) zur Schätzung der Skalierungsfaktoren, was eine hohe Genauigkeit gewährleistet, aber einen erheblichen Rechenaufwand verursacht.
2. Fast Mode (Schneller Modus): Verwendet die Cauchy-Schwarz-Ungleichung zur Schätzung der Skalierungsfaktoren ohne die zusätzliche GEMM, was einen höheren Durchsatz bietet, aber historisch gesehen eine geringere Genauigkeit aufweist.

Dieses Paper identifiziert einen kritischen Fehler in der bestehenden Skaliierungsformel des Fast Mode: Sie mangelt an Skaleninvarianz (scale invariance). Wenn die Eingangsmatrizen mit einer Skalar-Konstante $\alpha$ multipliziert werden, ändert sich die effektive Bitbreite der resultierenden skalierten Integer-Matrizen unvorhersehbar.

• Bei großen $\alpha$ sinkt die effektive Bitbreite, was zu einer Verschlechterung der Genauigkeit führt.
• Bei kleinen $\alpha$ kann das skalierte Integer-Produkt den CRT-Schwellenwert ($P/2$) überschreiten, was zum vollständigen Scheitern der Rekonstruktion führt.
  Numerische Experimente mit Einsen-Matrizen und Zufallsmatrizen zeigten, dass der Fast Mode für bestimmte Skalarbereiche nicht in der Lage ist, das korrekte Ergebnis zu rekonstruieren, während der Accurate Mode stabil bleibt.

Methodik
Die Autoren schlagen eine revidierte Skalierungsformel für den Fast Mode vor, die direkt aus der CRT-Eindeftigkeitsbedingung mittels der Cauchy-Schwarz-Ungleichung abgeleitet wurde.

1. Theoretische Herleitung:
   Die Autoren formulieren die Bedingung für die Skalierungsfaktoren $\mu_i$ und $\nu_j$ neu. Anstatt der bisherigen Heuristik definieren sie eine hinreichende Bedingung, bei der gilt:
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   Durch Lösen nach $\mu_i$ und $\nu_j$ unter Verwendung von Basis-2-Logarithmen und Berücksichtigung der absoluten Fehlerschranken der log2f-Funktion, die von GPU-Mathematik-APIs bereitgestellt wird, leiten sie neue Formeln ab:
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   wobei $P'_{prop}$ eine vorberechnete Konstante und $\widehat{\log_2}$ ein beschränkter Näherungswert des Logarithmus ist.

2. Beweis der Skaleninvarianz:
   Die Autoren demonstrieren analytisch, dass die neue Formel skaleninvariant ist. Wenn die Eingangsmatrix $A$ mit $\alpha = 2^s$ skaliert wird, kürzt sich der Term $\log_2(\alpha)$ in der Berechnung der effektiven Bitbreite der skalierten Integer-Matrix heraus. Folglich bleibt die Bitbreite unabhängig von der Größenordnung der Eingabe konstant, was die im ursprünglichen Fast Mode beobachtete Verschlechterung der Genauigkeit und das Versagen des CRT verhindert.

3. Implementierung:
   Die vorgeschlagene Methode (OS II-prop) ersetzt die Skalierungsformel im Fast Mode der GEMMul8-Bibliothek. Sie behält die Komplexität von $O(mk)$ und $O(kn)$ des ursprünglichen Fast Mode bei, führt keinen zusätzlichen Overhead gegenüber dem ursprünglichen Fast Mode ein und bleibt schneller als der Accurate Mode (der eine zusätzliche GEMM erfordert).

Zentrale Beiträge

• Identifizierung der Einschränkung: Das Paper zeigt auf, dass die bestehende Skalierungsformel des Fast Mode keine Skaleninvarianz besitzt, was zu Genauigkeitsverlusten und zum Scheitern der CRT-Rekonstruktion bei spezifischen Eingangsskalierungen führt.
• Revidierte Formel: Eine neue Skalierungsformel wird vorgeschlagen, die Skaleninvarianz garantiert und sicherstellt, dass die CRT-Eindeftigkeitsbedingung stets erfüllt ist.
• Null-Overhead: Die neue Formel führt keine zusätzlichen Rechenkosten gegenüber dem ursprünglichen Fast Mode ein.
• Umfassende Evaluierung: Die Methode wird auf einer NVIDIA GH200 GPU sowohl für Doppelpräzisions- (DGEMM) als auch für Einfachpräzisions-Matrixmultiplikationen (SGEMM) evaluiert.

Ergebnisse
Die Experimente wurden mit Zufallsmatrizen mit variierenden Elementgrößen-Spreads ($\phi$) und Matrizendimensionen ($k \in \{1024, 16384\}$) durchgeführt.

• Genauigkeit: Die vorgeschlagene Methode (OS II-prop) erreicht eine Genauigkeit, die über alle Werte von $\phi$ hinweg mit dem Accurate Mode (OS II-accu) vergleichbar ist. Im Gegensatz dazu zeigte der ursprüngliche Fast Mode (OS II-fast) eine signifikante Verschlechterung der Genauigkeit und ein Versagen bei großen $\phi$-Werten (ill-conditioned Fälle).
• Durchsatz: OS II-prop behält einen Durchsatz bei, der mit dem ursprünglichen Fast Mode vergleichbar ist. Während OS II-prop in spezifischen Szenarien mit geringen Moduli-Anzahlen etwas langsamer als der ursprüngliche Fast Mode ist (aufgrund weniger Null-Elemente in den skalierten Integern), ist es im Durchsatz signifikant schneller als der Accurate Mode.
• Trade-off: In der Analyse des Genauigkeits-Durchsatz-Verhältnisses dominiert OS II-prop gegenüber OS II-accu (höherer Durchsatz bei gleicher Genauigkeit) und überwindet die Genauigkeitsbeschränkungen von OS II-fast. Es ist die einzige Methode, die in der Lage ist, cuBLAS über einen weiten Bereich von Bedingungen hinweg sowohl in der Genauigkeit als auch im Durchsatz gleichzeitig zu übertreffen.

Bedeutung
Das Paper behauptet, dass die vorgeschlagene Methode die fundamentale Genauigkeitsbeschränkung des Fast Mode im Ozaki-Schema II löst, ohne dessen Leistungsvorteil zu opfern. Durch die Gewährleistung der Skaleninvarianz garantiert die Methode, dass die CRT-Rekonstruktion unabhängig von der Skalierung der Eingabe erfolgreich ist, was ein besseres Gleichgewicht zwischen Genauigkeit und Leistung bietet. Dies ermöglicht es Anwendern, hochpräzise Ergebnisse auf Low-Precision-Hardware mit der Geschwindigkeit des Fast Mode und der Zuverlässigkeit des Accurate Mode zu erzielen, wodurch die Durchsatzbeschränkungen des Accurate Mode und die Genauigkeitsbeschränkungen des ursprünglichen Fast Mode effektiv überwunden werden.

Die Autoren merken an, dass sich die aktuelle Arbeit auf INT8-Matrix-Engines für DGEMM und SGEMM konzentriert, der Ansatz jedoch auch für FP8-Engines und komplexe Matrixmultiplikationen (ZGEMM/CGEMM) anwendbar ist, welche für zukünftige Arbeiten vorgesehen sind.