改进Shampoo优化器速度,让大模型训练更快更高效。
DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers
- 将预条件块堆叠成3D张量,提升GPU利用率。
- 用Newton-DB和切比雪夫多项式加速矩阵开方计算,提速5.6倍。
- 适合需要快速收敛的大模型训练场景。
Shampoo是领先的近似二阶优化器之一:其变体曾赢得MLCommons AlgoPerf竞赛,并能生成激活异常值更少、更易压缩的模型。然而,当前应用Shampoo因内部操作昂贵而带来显著计算延迟。本文提出DASH(分布式加速Shampoo),通过两项核心技术解决此问题:首先,将预条件块堆叠为3D张量,大幅提升GPU利用率;其次,引入Newton-DB迭代与切比雪夫多项式逼近法,作为计算Shampoo所需逆矩阵平方根的新方法。同时,我们首次深入分析了矩阵缩放对Shampoo收敛性的影响。实际测试中,我们的GPU感知实现相比优化良好的分布式Shampoo,最快可提升5.6倍优化步速,且Newton-DB在所有测试方法中达到最低验证困惑度。代码已开源。
原文摘要 · Abstract (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 \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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。