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

arXiv:2610.02199v1 公告类型：新 摘要：对大型语言模型（LLM）进行全参数微调会产生大量优化器状态内存开销，限制了现代 GPU 能够容纳的模型规模。现有方法要么压缩优化器状态，要么放弃一阶梯度，要么在保留稠密状态的同时改变更新几何结构。近期提出的 Muon 优化器通过矩阵值更新降低了优化器内存占用。然而，其几何结构不同于 AdamW，在对经过 AdamW 预训练的模型进行微调时可能导致性能下降。为在 LLM 微调中降低优化器内存占用且不牺牲准确率或计算效率，我们提出三值绝对最大值按列单稀疏优化器（Ternary Absolute-max Column-wise One-sparse optimizer），简称 TACO。TACO 遵循 Muon 基于算子范数的最速下降观点，并进一步推进这一几何路径：对于二维权重矩阵，TACO 在经过维度归一化的 $1\to1$ 算子范数下，通过选择每列中绝对值最大的元素的符号，计算精确的最速下降方向。该方法保留了一阶梯度，同时使优化器状态内存几乎可以忽略不计。实际使用的 TACO 优化器每列仅保留少量低精度梯度分量；在 OPT-13B 上，相较于 AdamW8bit，它将持久化优化器状态减少了 174 倍（从 27.7 GB 降至 0.16 GB），并将训练峰值内存减少了 2.9 倍（从 80.6 GB 降至 27.5 GB），同时实现了相当的准确率和运行时间。TACO 还支持在单张 80 GB H100 GPU 上，对来自多个模型系列、用于多种任务的 30B–32B 参数模型进行全参数微调。