SAGE：面向内存高效 LLM 优化的符号自适应梯度

AdamW 优化器虽然是 LLM 预训练的标准选择，但是一个关键的内存瓶颈，消耗的优化状态等价于模型大小的两倍。虽然像 SinkGD 这样的轻量级优化器试图解决此问题，但我们发现嵌入层的困境：这些方法无法处理嵌入层中稀疏且高方差的梯度，迫使采用混合设计以恢复 AdamW，从而部分抵消了内存收益。我们提出 SAGE (Sign Adaptive GradiEnt)，一种新型优化器，通过替换混合结构中的 AdamW 来解决此困境。SAGE 将 Lion 风格的更新方向与一种新的、内存高效的 $O(d)$ 自适应尺度相结合。该尺度充当 "安全调制器"，可被证明受限于 1.0，能够更有效地抑制高方差维度。这种更好的稳定性使 SAGE 能够实现更好的收敛性能。在 Llama 参数量最高达 13 亿的模型上，我们基于 SAGE 的混合方案实现了新的 perplexity 状态最佳成绩，超越了所有基线，包括 SinkGD 混合方案，同时显著减少了优化状态的内存占用。