SOAP、Muon 及其他：推动 LLM 预训练量表

Muon 和 SOAP 等高阶优化器提供比 AdamW 更快的收敛速度，但它们的计算成本和数值稳定性挑战限制了大规模采用。在这项工作中，我们采用并增强了预处理梯度方法，以克服大规模 LLM 预训练的实际挑战。我们首先识别大批量下 SOAP 的不稳定性，并提出算法修改，包括每步 QR 正交化和改进的预处理策略，以消除损失峰值并在这些机制中实现稳定的训练。然后，我们使用更新 RMS 匹配对 SOAP、Muon 和 AdamW 进行统一的实证研究，以确保优化器之间公平的学习率转移。作为此分析的一部分，我们根据经验评估了 Muon的正交化质量。我们对在数万亿个词元上训练的数十亿参数模型进行的实验表明，在我们测试的规模上，SOAP 和 Muon 始终优于 AdamW。值得注意的是，在用于下一个词元预测的批量大小高达 100M 词元的情况下，这些优化器保持了训练稳定性和质量，而 AdamW 却下降了。为了实现大规模的高效训练，我们引入了与 Megatron-LM 兼容的分层分布式优化器。我们的实现平衡了内存并隐藏了通信，同时避免了优化器计算的近似，从而保留了它们的收敛优势。此外，我们还确定并构建特定的系统级改进，以进一步加速我们的分层实施。为了支持研究社区，我们发布了一个包含新兴优化算法的代码库：https://github.com/NVIDIA-NeMo/Emerging-Optimizers