DASH: Shampoo más rápido mediante preacondicionamiento de bloques por lotes y solucionadores eficientes de raíz inversa
DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers
February 2, 2026
Autores: Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh
cs.AI
Resumen
Shampoo es uno de los principales optimizadores aproximados de segundo orden: una variante del mismo ganó la competición MLCommons AlgoPerf, y se ha demostrado que produce modelos con menos valores atípicos en las activaciones que son más fáciles de comprimir. Sin embargo, aplicar Shampoo actualmente conlleva el coste de una ralentización computacional significativa, debido a sus costosas operaciones internas. En este artículo, damos un paso importante para abordar esta deficiencia proponiendo \method (por Distributed Accelerated SHampoo), una implementación más rápida de Shampoo Distribuido basada en dos técnicas nuevas principales: En primer lugar, demostramos que los bloques del precondicionador pueden apilarse en tensores 3D para mejorar significativamente la utilización de la GPU; en segundo lugar, introducimos la iteración Newton-DB y las aproximaciones polinómicas de Chebyshev como enfoques novedosos y más rápidos para calcular las raíces cuadradas inversas de matrices requeridas por Shampoo. Junto con estas contribuciones algorítmicas, proporcionamos un primer análisis en profundidad de cómo el escalado de matrices afecta críticamente a la convergencia de Shampoo. En el plano práctico, nuestra implementación optimizada para GPU logra pasos del optimizador hasta 4.83 veces más rápidos en comparación con la bien optimizada versión Distribuida de Shampoo, mientras que Newton-DB alcanza la menor perplejidad de validación por iteración entre todos los métodos probados. Nuestro código está disponible en https://github.com/IST-DASLab/DASH.
English
Shampoo is one of the leading approximate second-order optimizers: a variant of it has won the MLCommons AlgoPerf competition, and it has been shown to produce models with lower activation outliers that are easier to compress. Yet, applying Shampoo currently comes at the cost of significant computational slowdown, due to its expensive internal operations. In this paper, we take a significant step to address this shortcoming by proposing \method (for Distributed Accelerated SHampoo), a faster implementation of Distributed Shampoo based on two main new techniques: First, we show that preconditioner blocks can be stacked into 3D tensors to significantly improve GPU utilization; second, we introduce the Newton-DB iteration and the Chebyshev polynomial approximations as novel and faster approaches for computing the inverse matrix roots required by Shampoo. Along with these algorithmic contributions, we provide a first in-depth analysis of how matrix scaling critically affects Shampoo convergence. On the practical side, our GPU-aware implementation achieves up to 4.83times faster optimizer steps compared to the well-optimized Distributed Shampoo, while Newton-DB attains the lowest validation perplexity per iteration among all tested methods. Our code is available at https://github.com/IST-DASLab/DASH.