DASH: Ускорение Shampoo за счет пакетного блочного предобуславливания и эффективных решателей обратного корня

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

February 2, 2026

Авторы: Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

cs.AI

Аннотация

Shampoo является одним из ведущих приближенных оптимизаторов второго порядка: его вариант победил в соревновании MLCommons AlgoPerf, и было показано, что он порождает модели с меньшим количеством выбросов в активациях, которые легче сжимать. Тем не менее, применение Shampoo в настоящее время сопряжено со значительным замедлением вычислений из-за дорогостоящих внутренних операций. В данной статье мы делаем значительный шаг к устранению этого недостатка, предлагая метод DASH (Distributed Accelerated SHampoo) — более быструю реализацию распределенного алгоритма Shampoo, основанную на двух основных новых методах. Во-первых, мы показываем, что блоки предобуславливателей можно объединять в 3D-тензоры для значительного повышения загрузки GPU; во-вторых, мы представляем итерацию Ньютона-Дэниэля-Бе́клемишева (Newton-DB) и аппроксимации полиномами Чебышёва как новые и более быстрые подходы для вычисления обратных матричных корней, требуемых Shampoo. Наряду с этими алгоритмическими достижениями, мы предоставляем первый глубокий анализ того, как масштабирование матриц критически влияет на сходимость Shampoo. С практической стороны, наша реализация, оптимизированная для GPU, демонстрирует до 4.83-кратного ускорения шагов оптимизатора по сравнению с хорошо оптимизированным распределенным Shampoo, в то время как метод Newton-DB достигает наименьшей перплексии на валидационной выборке на итерацию среди всех протестированных методов. Наш код доступен по адресу 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.