🤖 machine learning

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

本文介绍了 DASH，这是一种显著提升了速度的 Distributed Shampoo 优化器实现，它利用 3D 张量堆叠来提高 GPU 利用率，并采用新型的反平方根求解器（Newton-DB 和 Chebyshev 近似），在保持或提高收敛质量的同时，实现了高达 5.6 倍的加速。

原作者： Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

发布于 2026-06-26

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在试图教一个巨大的机器人学习一种新的语言。为了做到这一点，你需要一个“优化器”——一个聪明的教练，在每次课程后调整机器人的大脑（其参数），以最大限度地减少错误。

多年来，最受欢迎的教练是 Adam，这种方法很快，但有点“懒惰”。它只关注每个单独的神经元需要改变多少，却忽略了神经元是如何协同工作的。

Shampoo 登场了，它是一位聪明得多的教练。Shampoo 不再一个接一个地观察神经元，而是观察它们如何成组进行交互。这带来了更好的学习效果，并且让模型在后期更容易被压缩。然而，它有一个弊端：Shampoo 极其缓慢。它就像一位天才教练，因为花了太多时间计算完美的招式，导致机器人几乎没怎么得到练习。

这篇论文介绍了 DASH（分布式加速型 Shampoo），这是一个全新的系统，它能让这位天才教练以短跑运动员的速度运行，同时又不失其智慧。以下是他们实现这一目标的方法，使用了简单的类比：

1. “堆叠”技巧（分批块预处理）

问题：
想象你有一个庞大的图书库（数据）需要整理。旧的方法（分布式 Shampoo）是拿一本书，整理好，放回去，再拿下一本书，整理，如此循环。即使你有 1,000 名工人（GPU），他们大部分时间也都在等待前一个人完成。这是低效的。

DASH 的解决方案：
DASH 改变了工作流程。它不再是一个一个处理书籍，而是将它们堆叠成整齐、统一的 3D 塔状结构。现在，工人们可以抓起一整座塔，同时整理其中的所有书籍。

类比： 这就像收银员逐一扫描商品与传送带将一整箱杂货瞬间送入机器进行扫描之间的区别。
结果： 这种“堆叠”让计算机强大的图形芯片（GPU）能够满负荷运转，使优化步骤的速度提升了高达 5.6 倍。

2. “捷径”数学（高效逆平方根求解器）

问题：
为了完成工作，Shampoo 在每一步都必须解决一个非常困难的数学谜题：寻找一个巨大矩阵的“逆平方根”。解决旧方法的过程就像是在草堆里找针，必须逐一检查每一根草（这种方法被称为特征值分解）。它很精确，但慢得令人痛苦。

DSH 的解决方案：
作者引入了两个新的“捷径”来解决这个谜题：

Newton-DB： 一种聪明的迭代方法，通过每一次猜测不断接近答案，就像 GPS 在你驾驶时不断重新计算路线。
切比雪夫多项式（Chebyshev Polynomials）： 一种数学近似法，可以非常快速地预测答案。
类比： 与其在迷宫中走遍每一条路径来寻找出口（旧方法），这些新方法就像拥有一张显示大致方向的地图，让你能够向出口冲刺。

3. “尺子”校准（矩阵缩放）

问题：
在使用这些捷径时，如果数字太大或太小，数学运算可能会变得不稳定。旧方法使用的“尺子”（Frobenius 范数）太长了，拉伸了数字，导致捷径需要更多步才能收敛。这就像是用一把 100 英尺长的卷尺去测量一只小蚂蚁；精度丢失了。

DASH 的解决方案：
作者意识到他们需要一把更好的尺子。他们开发了一种称为**多幂迭代（Multi-Power-Iteration）**的技术。

类比： 与其用巨大的卷尺去猜测蚂蚁的长度，不如使用专门的高精度游标卡尺，使其完美契合蚂蚁。这确保了数学运算能够快速收敛，且不会因数值误差而崩溃。

4. 结果

论文在大型语言模型（拥有 9.53 亿参数的 Llama）上测试了 DASH。

速度： DASH 完成训练步骤中的“思考”部分比之前的最佳版本快了 5.6 倍。
质量： 尽管速度更快，但使用 DASH 训练的模型在学习方面同样出色，甚至略好。事实上，新的“Newton-DB”捷径产生的模型在所有测试方法中具有最低的误差率（困惑度）。

总结

可以将 Shampoo 想象成一台因为司机在走风景路线而困在交通堵塞中的法拉利引擎。DASH 是同一台法拉利引擎，但它现在拥有了：

更宽的高速公路（堆叠块）让它可以开得更快。
GPS 捷径（Newton-DB）来避开交通拥堵。
更好的导航工具（Multi-Power-Iteration）以确保它不会迷路。

其结果是，这个“聪明”的优化器（Shampoo）不再因为太慢而难以使用，使其成为了训练下一代 AI 模型的实用选择。

技术摘要：DASH —— 通过分块矩阵预处理与高效逆根求解器实现更快的 Shampoo 优化器

问题陈述

Shampoo 是一种领先的近似二阶优化器，以其能产生具有更低激活异常值和更高压缩率的模型而闻名。然而，其广泛应用受到显著计算开销的阻碍，主要是由于计算预处理矩阵（ $A^{-1/4}$ ）的逆矩阵根操作极其昂贵。虽然分布式 Shampoo（Distributed Shampoo）试图通过将预处理器拆分为 $B \times B$ 块来将复杂度从 $O(n^3)$ 降低到 $O(Bn^2)$ 来缓解这一问题，但现有的实现仍存在关键的效率缺陷：

串行处理： 标准实现按顺序处理数据块，未能充分利用现代 GPU 的并行性。
低效求解器： 默认依赖特征值分解（EVD）在 GPU 上难以并行化。其他的迭代方法（如 Coupled-Newton, CN）由于对数值稳定性和收敛速度的担忧，通常在默认情况下被禁用。
次优缩放： 当前方法通常使用 Frobenius 范数对矩阵进行缩放，这种方法计算成本虽低，但会将特征值推向零，从而需要更多的迭代次数才能达到目标精度。

方法论

本文提出了 DASH (Distributed Accelerated SHampoo)，这是一种旨在弥合理论效率与硬件能力之间差距的高性能实现。该方法论建立在两个主要支柱之上：

1. 分块矩阵预处理（系统级优化）

DASH 不再顺序处理预处理块，而是将独立的预处理块堆叠成 3D 张量。

3D 张量堆叠： 梯度块 $G_{ij}$ 被堆叠成形状为 $(N, B, B)$ 的矩阵，其中 $N$ 是总块数。
并行执行： 这使得优化器能够使用批处理矩阵乘法（bmm）原语，同时对整个块批次应用逆根程序（包括 EVD、CN 或新的迭代方法）。
GPU 利用率： 这种架构变化最大化了 GPU Tensor Core 的吞吐量，减少了内存碎片，并实现了对某些求解器的高效半精度（FP16）算术运算的使用。

2. 高效逆根求解器（算法级优化）

论文引入并分析了用于替代或增强标准 EVD 和 CN 的高级线性代数方法：

Newton-Denman-Beavers (NDB) 迭代： 一种用于 Shampoo 的新颖方法，可同时计算平方根和逆平方根。作者优化了第一次迭代以消除冗余的矩阵乘法，将开销降低至每步 3 次矩阵乘法（在实际应用中甚至更少）。
切比雪夫多项式逼近（Chebyshev Polynomial Approximations）： 利用 Clenshaw 算法来逼近矩阵幂，提供了一种替代的迭代求解器。
用于缩放的多幂迭代（Multi-Power-Iteration for Scaling）： 为了解决收敛问题，论文使用多幂迭代方法取代了标准的 Frobenius 范数缩放。
- 机制： 它通过一组并行处理的起始向量池（最多 32 个）来估计谱半径（ $\lambda_{max}$ ）。
- 益处： 这提供了比 Frobenius 范数更准确的缩放因子，确保特征值处于迭代求解器能快速收敛（接近 1）的区间内，从而减少所需的迭代次数并提高数值稳定性。

核心贡献

DASH 实现： 一个 GPU 高效的框架，通过将串行循环替换为并行 3D 张量操作，实现了比优化后的分布式 Shampoo 快达 5.6 倍 的优化器步进速度。
新型求解器： 引入了 Newton-DB (NDB) 迭代和 切比雪夫多项式 逼近作为 EVD 和 CN 的可行且更快的替代方案。论文证明了 NDB 比 CN 和标准 EVD 具有更低的验证困惑度（validation perplexity）。
收敛性分析： 详细的行为分析表明，标准的 Frobenius 范数缩放是次优的，往往会导致更慢的收敛。论文为 CN 与 NDB 之间截然不同的收敛行为提供了直观的解释。
多幂迭代： 一种高效的半精度幂迭代实现，能够稳健地估计谱半径，从而实现最优的矩阵缩放，使 NDB 和 CN 能以较少的步骤满足收敛标准。

实验结果

作者在拥有 953M 参数（嵌入大小 2048，序列长度 1024）并在 C4 数据集上训练的 Llama 模型上评估了 DASH。

速度：
- 根据所使用的逆根方法和精度不同，DASH 比分布式 Shampoo 减少了 4 倍至 5.6 倍 的优化器步进运行时间。
- 最显著的加速（5.6 倍）出现在将分布式 Shampoo 的 EVD (FP32) 与结合了 GPU 级堆叠的 DASH 的 CN (FP16) 进行对比时。
- 即使在频繁更新预处理器（ $f=1$ ）的情况下，DASH 仍保持了显著的速度优势。
准确度（验证困惑度）：
- DASH 在几乎所有稳定配置下都能达到与分布式 Shampoo 持平的效果。
- NDB 的验证困惑度始终低于 CN，并且通常优于 EVD。
- 使用带有 幂迭代缩放 的 NDB 得到了最低的验证困惑度（11.68），优于 Frobenius 缩放（11.76），验证了其理论分析。
精度：
- 该实现支持在不降低困惑度的情况下对 CN 和 EVD 使用 FP16，从而进一步缩短运行时间。
- 实验发现 NDB 在 FP16/BF16 格式下存在数值不稳定性的问题，但作者指出这是未来的研究方向。

意义与主张

论文声称 DASH 显著降低了使用像 Shampoo 这样的二阶优化器的计算门槛，使得高质量的预处理在处理大规模深度学习任务时变得切实可行。

桥接理论与硬件： 通过将算法设计（迭代求解器、矩阵缩量）与硬件能力（Tensor Core、批处理操作）相匹配，DASH 允许研究人员在没有过高运行时间负担的情况下研究优化器的特性。
数值稳定性 vs. 速度： 本研究挑战了“迭代方法本质上不如 EVD 稳定或更慢”的假设。通过引入更好的缩放（多幂迭代）和优化的迭代（NDB），作者证明了迭代方法可以实现更优的收敛性和更低的困惑度，同时速度更快。
未来方向： 作者建议，NDB 的卓越表现为基于块条件数（block condition numbers）的动态求解器选择开辟了道路，同时也为如何利用随机或误差校正方法在低精度格式下稳定 NDB 提供了可能。

论文总结道，Shampoo 的主要瓶颈并非数学复杂性，而是破碎的执行过程和数值不稳定性，而 DASH 有效地解决了这些问题。