Resumen Técnico: DASH – Shampoo más rápido mediante el preacondicionamiento de bloques por lotes y solucionadores de raíz inversa eficientes

Planteamiento del problema

Shampoo es un optimizador de segundo orden aproximado líder, conocido por producir modelos con menores valores atípicos en las activaciones y una mejor compresibilidad. Sin embargo, su adopción generalizada se ve obstaculizada por una sobrecarga computacional significativa, principalmente debido a la costosa operación de calcular raíces de matrices inversas ($A^{-1/4}$) para las matrices de preacondicionamiento. Aunque el Shampoo Distribuido intenta mitigar esto dividiendo los preacondicionadores en bloques de $B \times B$ para reducir la complejidad de $O(n^3)$ a $O(Bn^2)$, las implementaciones existentes sufren de brechas de eficiencia críticas:

1. Procesamiento Secuencial: Las implementaciones estándar procesan los bloques de forma secuencial, fallando al no aprovechar plenamente el paralelismo de las GPU modernas.
2. Solucionadores Ineficientes: La dependencia por defecto de la Descomposición de Autovalores (EVD) es difícil de paralelizar en GPU. Los métodos iterativos alternativos como Coupled-Newton (CN) suelen estar desactivados por defecto debido a preocupaciones sobre la estabilidad numérica y la velocidad de convergencia.
3. Escalado Subóptimo: Los métodos actuales suelen escalar las matrices utilizando la norma de Frobenius, lo cual es computacionalmente barato pero puede empujar los autovalores hacia cero, requiriendo significativamente más iteraciones para converger a una precisión objetivo.

Metodología

El artículo propone DASH (Distributed Accelerated SHampoo), una implementación de alto rendimiento diseñada para cerrar la brecha entre la eficiencia teórica y las capacidades del hardware. La metodología se basa en dos pilares principales:

1. Preacondicionamiento de Bloques por Lotes (Optimización a Nivel de Sistema)

En lugar de procesar los bloques del preacondicionador de forma secuencial, DASH apila bloques independientes en tensores 3D.

• Apilamiento de Tensores 3D: Los bloques de gradiente $G_{ij}$ se apilan en una matriz de forma $(N, B, B)$, donde $N$ es el número total de bloques.
• Ejecución Paralela: Esto permite que el optimizador aplique procedimientos de raíz inversa (EVD, CN o nuevos métodos iterativos) a todo el lote de bloques simultáneamente utilizando primitivas de multiplicación de matrices por lotes (bmm).
• Utilización de GPU: Este cambio arquitectónico maximiza el rendimiento de los Tensor Cores de la GPU, reduciendo la fragmentación de memoria y permitiendo el uso eficiente de la aritmética de media precisión (FP16) para ciertos solucionadores.

2. Solucionadores de Raíz Inversa Eficientes (Optimización Algorítmica)

El artículo introduce y analiza enfoques de álgebra lineal avanzada para reemplazar o complementar la EVD y CN estándar:

• Iteración Newton-Denman-Beavers (NDB): Un enfoque novedoso para Shampoo que calcula tanto la raíz cuadrada como la raíz cuadrada inversa. Los autores optimizan la primera iteración para eliminar multiplicaciones de matrices redundantes, reduciendo la sobrecarga a 3 multiplicaciones de matrices por paso (efectivamente menos en la práctica).
• Aproximaciones de Polinomios de Chebyshev: Utilizando el algoritmo de Clenshaw para aproximar potencias de matrices, ofreciendo un solucionador iterativo alternativo.
• Multi-Power-Iteration para el Escalado: Para abordar los problemas de convergencia, el artículo reemplaza el escalado estándar de la norma de Frobenius con un método de Multi-Power-Iteration.
  • Mecanismo: Estima el radio espectral ($\lambda_{max}$) utilizando un conjunto de vectores iniciales (hasta 32) procesados en paralelo mediante Power-Iteration.
  • Beneficio: Proporciona un factor de escala más preciso que la norma de Frobenius, asegurando que los autovalores permanezcan en un régimen donde los solucionadores iterativos convergen rápidamente (cerca de 1), reduciendo así el número de iteraciones requeridas y mejorando la estabilidad numérica.

Contribuciones Clave

1. Implementación DASH: Un marco de trabajo eficiente para GPU que reemplaza los bucles secuenciales con operaciones de tensores 3D paralelas, logrando hasta 5.6× más rápido los pasos del optimizador en comparación con el Shampoo Distribuido bien optimizado.
2. Nuevos Solucionadores: La introducción de las iteraciones Newton-DB (NDB) y las aproximaciones de polinomios de Chebyshev como alternativas viables y más rápidas a EVD y CN. El artículo demuestra que NDB logra una menor perplejidad de validación que tanto Coupled-Newton como la EVD estándar.
3. Análisis de Convergencia: Un análisis detallado del comportamiento que muestra que el escalado estándar de la norma de Frobenius es subóptimo para los métodos iterativos, lo que a menudo conduce a una convergencia más lenta. El artículo proporciona explicaciones intuitivas para los distintos comportamientos de convergencia de CN frente a NDB.
4. Multi-Power-Iteration: Una implementación eficiente de media precisión de Power-Iteration que estima robustamente los radios espectrales, permitiendo un escalado de matrices óptimo que satisface los criterios de convergencia para NDB y CN con menos pasos.

Resultados Experimentales

Los autores evaluaron DASH en un modelo Llama con 953 millones de parámetros (tamaño de embedding 2048, longitud de secuencia 1024) entrenado en el conjunto de datos C4.

• Velocidad:
  • DASH reduce el tiempo de ejecución del paso del optimizador de 4× a 5.6× en comparación con el Shampoo Distribuido, dependiendo del método de raíz inversa y la precisión utilizada.
  • La aceleración más significativa (5.6×) se observó al comparar la EVD (FP32) de Distributed Shampoo contra la CN (FP16) de DASH con el apilamiento a nivel de GPU.
  • Incluso con actualizaciones frecuentes del preacondicionador ($f=1$), DASH mantiene ventajas sustanciales de velocidad.
• Precisión (Perplejidad de Validación):
  • DASH iguala al Shampoo Distribuido en casi todas las configuraciones estables.
  • NDB ofrece consistentemente una menor perplejidad de validación que CN y a menudo supera a EVD.
  • El uso de escalado por Power-Iteration con NDB resultó en la menor perplejidad de validación (11.68) en comparación con el escalado de Frobenius (11.76), validando el análisis teórico.
• Precisión:
  • La implementación soporta FP16 para CN y EVD sin degradar la perplejidad, reduciendo aún más el tiempo de ejecución.
  • Se encontró que NDB en FP16/BF16 fue numéricamente inestable en sus experimentos, aunque los autores señalan esto como un área para trabajos futuros.

Significado y Reivindicaciones

El artículo afirma que DASH reduce significativamente la barrera computacional para el uso de optimizadores de segundo orden como Shampoo, haciendo práctico el uso de un preacondicionamiento de alta calidad a gran escala.

• Puente entre la Teoría y el Hardware: Al alinear el diseño algorítmico (solucionadores iterativos, escalado de matrices) con las capacidades del hardware (Tensor Cores, operaciones por lotes), DASH permite a los investigadores investigar las propiedades de los optimizadores sin tiempos de ejecución prohibitivos.
• Estabilidad Numérica vs. Velocidad: El trabajo desafía la suposición de que los métodos iterativos son inherentemente menos estables o más lentos que EVD. Al introducir un mejor escalado (Multi-Power-Iteration) y las iteraciones optimizadas (NDB), los autores demuestran que los métodos iterativos pueden lograr una convergencia superior y una menor perplejidad siendo significativamente más rápidos.
• Direcciones Futuras: Los autores sugieren que el rendimiento superior de Newton-DB abre vías para la selección dinámica de solucionadores basada en los números de condición de los bloques y la posibilidad de estabilizar NDB en formatos de menor precisión mediante métodos estocásticos o de corrección de errores.

El artículo concluye que el cuello de botella principal de Shampoo no era la complejidad matemática, sino la ejecución fragmentada y la inestabilidad numérica, ambos aspectos que DASH aborda de manera efectiva.