DASH : Shampoo accéléré par préconditionnement par blocs par lots et solveurs efficaces d'inverse-racine

DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers

February 2, 2026

papers.authors: Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

cs.AI

papers.abstract

Shampoo est l'un des optimiseurs approximatifs du second ordre les plus performants : une variante a remporté la compétition MLCommons AlgoPerf, et il a été démontré qu'il produit des modèles avec moins de valeurs aberrantes dans les activations, ce qui les rend plus faciles à compresser. Cependant, l'application de Shampoo se fait actuellement au prix d'un ralentissement computationnel significatif, en raison de ses opérations internes coûteuses. Dans cet article, nous franchissons une étape importante pour remédier à cet inconvénient en proposant \method (pour Distributed Accelerated SHampoo), une implémentation plus rapide de Shampoo Distribué basée sur deux nouvelles techniques principales : Premièrement, nous montrons que les blocs de préconditionneurs peuvent être empilés en tenseurs 3D pour améliorer considérablement l'utilisation du GPU ; deuxièmement, nous introduisons l'itération de Newton-DB et les approximations par polynômes de Tchebychev comme approches novatrices et plus rapides pour calculer les racines carrées inverses des matrices requises par Shampoo. Parallèlement à ces contributions algorithmiques, nous fournissons une première analyse approfondie de la manière dont la mise à l'échelle des matrices affecte de façon critique la convergence de Shampoo. Sur le plan pratique, notre implémentation optimisée pour les GPU permet d'atteindre jusqu'à 4.83 fois des étapes d'optimisation plus rapides par rapport à la version bien optimisée de Shampoo Distribué, tandis que Newton-DB atteint la plus faible perplexité de validation par itération parmi toutes les méthodes testées. Notre code est disponible à l'adresse 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.