soap

一篇新的NVIDIA论文提出，像Muon和SOAP这样的高阶优化器，可以作为大规模LLM预训练中AdamW的更高效替代方案。

本文改进了高阶优化器 SOAP 和 Muon，用于大规模 LLM 预训练，解决了大批量训练中的不稳定性问题，并引入了一种与 Megatron-LM 兼容的层式分布式优化器。实验表明，在数十亿参数规模下，它们始终优于 AdamW。

本文提出了一种对基于Shampoo的优化方法（如KL-Shampoo和SOAP）中预条件器的重新参数化，以支持BFloat16存储，并通过在子空间中仅通过QR分解更新部分基来减少计算开销，从而使这些方法在内存和时间上更高效。