DASH: 通过批量块预处理和高效逆根求解器加速 Shampoo
DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers
AI总结：
提出 DASH 方法，通过将预处理块堆叠为3D张量提升GPU利用率，并引入 Newton-DB 迭代和 Chebyshev 多项式近似加速逆矩阵根计算，实现高达5.6倍的优化器加速。
摘要：
Shampoo 是领先的近似二阶优化器之一：其变体赢得了 MLCommons AlgoPerf 竞赛，并且已被证明能产生具有较低激活异常值且更易压缩的模型。然而，由于 Shampoo 内部操作昂贵，目前应用它会导致显著的计算速度下降。在本文中，我们通过提出 \method（代表 \textbf{D}istributed \textbf{A}ccelerated \textbf{SH}ampoo）迈出了解决这一缺陷的重要一步，这是一种基于两项主要新技术的分布式 Shampoo 的更快实现：首先，我们展示了预处理块可以堆叠成 3D 张量，以显著提高 GPU 利用率；其次，我们引入了 Newton-DB 迭代和 Chebyshev 多项式近似，作为计算 Shampoo 所需的逆矩阵根的新颖且更快的方法。除了这些算法贡献外，我们还首次深入分析了矩阵缩放如何关键地影响 Shampoo 收敛。在实际方面，与优化良好的分布式 Shampoo 相比，我们的 GPU 感知实现实现了高达 5.6 倍的优化器步骤加速，而 Newton-DB 在所有测试方法中实现了每次迭代最低的验证困惑度。我们的代码可在以下网址获取：this https URL。
英文摘要：
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 \textbf{D}istributed \textbf{A}ccelerated \textbf{SH}ampoo), 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 $5.6\times$ 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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
