MuonQ：通过定向保真度优化增强低位 Muon 量化

Muon 优化器已成为 Adam 训练 大语言模型 的一个引人注目的替代方案，通过梯度正交化实现了显着的计算节省。然而，Muon 的优化器状态对量化误差更敏感：因为正交化丢弃了奇异值的大小并仅保留方向信息，所以即使奇异向量方向上的小量化误差也会在更新中被放大。在这项工作中，我们提出了 MuonQ，一种基于方向保真度优化原理的低位 Muon 训练框架。首先，我们应用预量化归一化，以便每一步都会引入相同幅度的量化误差，从而防止累积误差发展为首选方向。其次，我们引入了一种结构分解，通过幂迭代单独量化主要奇异分量，确保量化误差仅扰动奇异值幅度而不是旋转奇异向量方向。第三，我们采用 $μ$ 律压扩量化，将更高分辨率分配给密集动量值，将量化目标从异常值保留转变为密集区域可区分性。这些技术共同实现了 Muon 优化器状态的稳定 4 位量化。针对 GPT 型和 LLaMA 型模型的 预训练 实验表明，4 位精度的 MuonQ 可以恢复大部分全精度 Muon 的训练损失和下游精度，同时将优化器状态内存减少高达 7.3$\times$，在精度与内存权衡方面提供了有利点。我们的代码可在 https://github.com/YupengSu/MuonQ 获取。