Improved Scaling for Fast Mode of Ozaki Scheme II

Sintesi Tecnica: Miglioramento della Scalabilità per la Fast Mode dello Schema Ozaki II

Enunciato del Problema
Lo Schema Ozaki II è un metodo per emulare la moltiplicazione di matrici ad alta precisione (es. FP64/FP32) utilizzando hardware a bassa precisione (es. INT8 Tensor Cores) basato sul Teorema Cinese del Resto (CRT). Lo schema richiede di scalare le matrici di input ad alta precisione in matrici intere tali che il prodotto di questi interi, quando ricostruito tramite CRT, non superi una soglia specifica ($P/2$). Se questa soglia viene superata, il recupero tramite CRT fallisce, portando a una perdita catastrofica di accuratezza.

Lo schema offre due modalità per determinare i fattori di scala ($\mu$ e $\nu$):

1. Accurate Mode: Utilizza una moltiplicazione di matrici (GEMM) ausiliaria in INT8 per stimare i fattori di scala, garantendo un'elevata accuratezza ma comportando un carico computazionale significativo.
2. Fast Mode: Utilizza la disuguaglianza di Cauchy–Schwarz per stimare i fattori di scala senza la GEMM ausiliaria, offrendo un throughput più elevato ma storicamente un'accuratezza inferiore.

Questo articolo identifica un difetto critico nella formula di scaling esistente della Fast Mode: essa manca di invarianza di scala. Quando le matrici di input vengono moltiplicate per una costante scalare $\alpha$, la larghezza di bit effettiva delle matrici intere scalate risultanti cambia in modo imprevedibile.

• Per $\alpha$ grande, la larghezza di bit effettiva diminuisce, causando una degradazione dell'accuratezza.
• Per $\alpha$ piccolo, il prodotto intero scalato può superare la soglia del CRT ($P/2$), causando il fallimento totale del recupero.
  Esperimenti numerici su matrici composte da soli uno e su matrici casuali hanno dimostrato che, per determinati intervalli scalari, la Fast Mode non riesce a recuperare il risultato corretto, mentre l'Accurate Mode rimane stabile.

Metodologia
Gli autori propongono una formula di scaling rivista per la Fast Mode che è derivata direttamente dalla condizione di unicità del CRT tramite la disuguaglianza di Cauchy–Schwarz.

1. Derivazione Teorica:
   Gli autori riformulano la condizione per i fattori di scala $\mu_i$ e $\nu_j$. Invece della precedente euristica, definiscono una condizione sufficiente in cui:
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   Risolvendo per $\mu_i$ e $\nu_j$ utilizzando i logaritmi in base 2 e tenendo conto dei limiti di errore assoluto della funzione log2f fornita dalle API matematiche delle GPU, derivano nuove formule:
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   dove $P'_{prop}$ è una costante precomputata e $\widehat{\log_2}$ è un'approssimazione limitata del logaritmo.

2. Dimostrazione dell'Invarianza di Scala:
   Gli autori dimostrano analiticamente che la nuova formula è invariante di scala. Se la matrice di input $A$ è scalata per $\alpha = 2^s$, il termine $\log_2(\alpha)$ si cancella nel calcolo della larghezza di bit effettiva della matrice intera scalata. Di conseguenza, la larghezza di bit rimane costante indipendentemente dall'ampiezza dell'input, prevenendo la degradazione dell'accuratezza e il fallimento del CRT osservati nella Fast Mode originale.

3. Implementazione:
   Il metodo proposto (OS II-prop) sostituisce la formula di scaling nella Fast Mode della libreria GEMMul8. Mantiene la complessità computazionale $O(mk)$e$O(kn)$ della Fast Mode originale, non introducendo alcun overhead aggiuntivo rispetto alla Fast Mode originale e rimanendo più veloce della Accurate Mode (che richiede una GEMM extra).

Contributi Chiave

• Identificazione del Limite: Il documento dimostra che la formula di scaling della Fast Mode esistente manca di invarianza di scala, portando alla degradazione dell'accuratezza e al fallimento del recupero CRT per specifici scalamenti dell'input.
• Formula Rivista: Viene proposta una nuova formula di scaling che garantisce l'invarianza di scala e assicura che la condizione di unicità del CRT sia sempre soddisfatta.
• Zero Overhead: La nuova formula non introduce costi computazionali aggiuntivi rispetto alla Fast Mode originale.
• Valutazione Esaustiva: Il metodo è valutato su una GPU NVIDIA GH200 sia per la moltiplicazione di matrici a doppia precisione (DGEMM) che a singola precisione (SGEMM).

Risultati
Gli esperimenti sono stati condotti utilizzando matrici casuali con diverse dispersioni di ampiezza degli elementi ($\phi$) e dimensioni delle matrici ($k \in \{1024, 16384\}$).

• Accuratezza: Il metodo proposto (OS II-prop) raggiunge un'accuratezza comparabile alla Accurate Mode (OS II-accu) per tutti i valori di $\phi$. Al contrario, la Fast Mode originale (OS II-fast) ha mostrato una significativa degradazione dell'accuratezza e fallimenti per valori di $\phi$ elevati (casi mal condizionati).
• Throughput: OS II-prop mantiene un throughput comparabile alla Fast Mode originale. Sebbene OS II-prop sia leggermente più lento della Fast Mode originale in scenari specifici con piccoli conteggi di moduli (a causa di meno elementi zero negli interi scalati), supera significativamente la Accurate Mode in termini di throughput.
• Trade-off: Nell'analisi del compromesso accuratezza-throughput, OS II-prop domina OS II-accu (maggiore throughput per la stessa accuratezza) e supera i limiti di accuratezza di OS II-fast. È l'unico metodo capace di superare simultaneamente cuBLAS sia in accuratezza che in throughput per un ampio intervallo di condizioni.

Significato
Il documento afferma che il metodo proposto risolve il limite fondamentale di accuratezza della Fast Mode nello Schema Ozaki II senza sacrificarne il vantaggio prestazionale. Garantendo l'invarianza di scala, il metodo assicura che il recupero tramite CRT abbia successo indipendentemente dallo scaling dell'input, offrendo un equilibrio superiore tra accuratezza e prestazioni. Ciò consente agli utenti di ottenere risultati ad alta precisione su hardware a bassa precisione con la velocità della Fast Mode e l'affidabilità della Accurate Mode, superando efficacementmente i vincoli di throughput della Accurate Mode e i limiti di accuratezza della Fast Mode originale.

Gli autori osservano che, sebbene il presente lavoro si concentri sui motori di matrici INT8 per DGEMM e SGEMM, l'approccio è applicabile ai motori FP8 e alle moltiplicazioni di matrici complesse (ZGEMM/CGEMM), che sono lasciate per lavori futuri.