TACO：用于 LLM 微调的三元绝对最大列单稀疏优化器

大语言模型 (LLM) 的全参数微调会产生大量优化器状态内存开销，从而限制了适合现代 GPU 的模型大小。现有方法要么压缩优化器状态，放弃一阶梯度，要么在保留密集状态的同时更改更新几何结构。最近推出的 Muon 优化器通过矩阵值更新减少了优化器内存。尽管如此，它的几何结构与 AdamW 不同，并且在微调 AdamW 预训练模型时可能会导致性能下降。为了在不牺牲 LLM 微调精度或计算效率的情况下减少优化器内存，我们提出了三元绝对最大列单稀疏优化器（TACO），它遵循 Muon 的算子范数最速下降视图，但进一步采用了几何路线。 TACO 通过选择二维权重矩阵每列中最大幅度条目的符号，在维度归一化 $1\to1$ 算子范数下计算精确的最速下降方向。这保留了一阶梯度，同时使优化器状态内存几乎可以忽略不计。我们实用的 TACO 优化器每列仅维护一小组低精度梯度组件，相对于 AdamW8bit（从 27.7 GB 到 0.16 GB）减少了 $174\times$ 的持续优化器状态，并在 OPT-13B 上减少了 $2.9\times$ 的峰值训练内存（从 80.6 GB 到 27.5 GB），同时实现了相当的精度和运行时间。 TACO 还可以在单个 80 GB H100 GPU 上跨多个模型系列和任务对 30-32B 参数模型进行全参数微调。