Improved Scaling for Fast Mode of Ozaki Scheme II

Resumen Técnico: Mejora del Escalamiento para el Modo Rápido del Esquema de Ozaki II

Planteamiento del Problema
El Esquema de Ozaki II es un método para emular la multiplicación de matrices de alta precisión (p. ej., FP64/FP32) utilizando hardware de baja precisión de enteros (p. ej., Tensor Cores de INT8) basado en el Teorema del Resto Chino (CRT). El esquema requiere escalar las matrices de entrada de alta precisión en matrices de enteros tales que el producto de estos enteros, al ser reconstruido mediante CRT, no exceda un umbral específico ($P/2$). Si este umbral se excede, la recuperación de CRT falla, provocando una pérdida catastrófica de precisión.

El esquema ofrece dos modos para determinar los factores de escala ($\mu$ y $\nu$):

1. Modo Preciso (Accurate Mode): Utiliza una multiplicación de matrices (GEMM) de enteros INT8 auxiliar para estimar los factores de escala, lo que garantiza una alta precisión pero incurre en una sobrecarga computacional significativa.
2. Modo Rápido (Fast Mode): Utiliza la desigualdad de Cauchy–Schwarz para estimar los factores de escala sin la GEMM auxiliar, ofreciendo un mayor rendimiento (throughput) pero históricamente una menor precisión.

Este artículo identifica un fallo crítico en la fórmula de escalamiento existente del Modo Rápido: carece de invarianza de escala. Cuando las matrices de entrada se multiplican por una constante escalar $\alpha$, el ancho de bits efectivo de las matrices de enteros escaladas cambia de forma impredecible.

• Para $\alpha$ grande, el ancho de bits efectivo disminuye, causando degradación de la precisión.
• Para $\alpha$ pequeño, el producto de enteros escalados puede exceder el umbral de CRT ($P/2$), causando un fallo total en la recuperación.
  Experimentos numéricos con matrices de unos y matrices aleatorias demostraron que, para ciertos rangos de escalares, el Modo Rápido no logra recuperar el resultado correcto, mientras que el Modo Preciso permanece estable.

Metodología
Los autores proponen una fórmula de escalamiento revisada para el Modo Rápido que se deriva directamente de la condición de unicidad de CRT mediante la desigualdad de Cauchy–Schwarz.

1. Derivación Teórica:
   Los autores reformulan la condición para los factores de escala $\mu_i$ y $\nu_j$. En lugar de la heurística anterior, definen una condición suficiente donde:
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   Al resolver para $\mu_i$ y $\nu_j$ utilizando logaritmos en base 2 y contabilizando los límites de error absoluto de la función log2f proporcionada por las API matemáticas de GPU, derivan nuevas fórmulas:
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   donde $P'_{prop}$ es una constante precalculada y $\widehat{\log_2}$ es una aproximación acotada del logaritmo.

2. Prueba de Invarianza de Escala:
   Los autores demuestran analíticamente que la nueva fórmula es invariante de escala. Si la matriz de entrada $A$ se escala por $\alpha = 2^s$, el término $\log_2(\alpha)$ se cancela en el cálculo del ancho de bits efectivo de la matriz de enteros escalada. Consecuentemente, el ancho de bits permanece constante independientemente de la magnitud de la entrada, evitando la degradación de la precisión y el fallo de CRT observados en el Modo Rápido original.

3. Implementación:
   El método propuesto (OS II-prop) reemplaza la fórmula de escalamiento en el Modo Rápido de la biblioteca GEMMul8. Mantiene la complejidad computacional de $O(mk)$y$O(kn)$ del Modo Rápido original, sin introducir sobrecarga adicional en comparación con el Modo Rápido original y permaneciendo más rápido que el Modo Preciso (que requiere una GEMM adicional).

Contribuciones Clave

• Identificación de la Limitación: El artículo demuestra que la fórmula de escalamiento del Modo Rápido existente carece de invarianza de escala, lo que conduce a la degradación de la precisión y al fallo de la recuperación de CRT para escalamientos de entrada específicos.
• Fórmula Revisada: Se propone una nueva fórmula de escalamiento que garantiza la invarianza de escala y asegura que la condición de unicidad de CRT se cumpla siempre.
• Cero Sobrecarga: La nueva fórmula no introduce costo computacional adicional sobre el Modo Rápido original.
• Evaluación Exhaustiva: El método se evalúa en una GPU NVIDIA GH200 tanto para multiplicaciones de matrices de doble precisión (DGEMM) como de precisión simple (SGEMM).

Resultados
Los experimentos se realizaron utilizando matrices aleatorias con variaciones en la dispersión de la magnitud de los elementos ($\phi$) y dimensiones de matriz ($k \in \{1024, 16384\}$).

• Precisión: El método propuesto (OS II-prop) logra una precisión comparable al Modo Preciso (OS II-accu) en todos los valores de $\phi$. En contraste, el Modo Rápido original (OS II-fast) mostró una degradación significativa de la precisión y fallos para valores altos de $\phi$ (casos mal condicionados).
• Rendimiento (Throughput): OS II-prop mantiene un rendimiento comparable al Modo Rápido original. Aunque OS II-prop es ligeramente más lento que el Modo Rápido original en escenarios específicos con pocos módulos (debido a menos elementos cero en los enteros escalados), supera significativamente al Modo Preciso en rendimiento.
• Compromiso (Trade-off): En el análisis de compromiso entre precisión y rendimiento, OS II-prop domina a OS II-accu (mayor rendimiento para la misma precisión) y supera las limitaciones de precisión de OS II-fast. Es el único método capaz de superar simultáneamente a cuBLAS tanto en precisión como en rendimiento para un amplio rango de condiciones.

Significancia
El artículo afirma que el método propuesto resuelve la limitación fundamental de precisión del Modo Rápido en el Esquema de Ozaki II sin sacrificar su ventaja de rendimiento. Al garantizar la invarianza de escala, el método garantiza que la recuperación de CRT tenga éxito independientemente del escalamiento de la entrada, ofreciendo un equilibrio superior entre precisión y rendimiento. Esto permite a los usuarios obtener resultados de alta precisión en hardware de baja precisión con la velocidad del Modo Rápido y la fiabilidad del Modo Preciso, superando efectivamente las restricciones de rendimiento del Modo Preciso y las restricciones de precisión del Modo Rápido original.

Los autores señalan que, aunque el presente trabajo se centra en motores de matrices INT8 para DGEMM y SGEMM, el enfoque es aplicable a motores FP8 y multiplicaciones de matrices complejas (ZGEMM/CGEMM), lo cual queda para trabajos futuros.