TACO 优化器将 LLM 全参数微调的优化器状态内存压缩 174 倍，并在单张 80GB H100 上微调 30-32B 模型
相关研究与后续进展
核心概要
作者提出 TACO 优化器，在维度归一化的 1→1 算子范数下取二维权重矩阵每列最大幅值元素的符号作为精确最速下降方向，仅保留每列少量低精度梯度分量，在 OPT-13B 上把持久优化器状态相对 AdamW8bit 减少 174 倍（27.7 GB 降至 0.16 GB）、峰值训练内存减少 2.9 倍（80.6 GB 降至 27.5 GB），精度与运行时间相当，并支持在单张 80 GB H100 上对 30-32B 参数模型做全参数微调。
Figure 1: Performance–memory overview of TACO and other optimizers. Left: Pareto frontier memory-accuracy trade-off for fine-tuning OPT-13B on the SST-2 dataset on an 80 GB H100 GPU. TACO lies at the low-memory end of the Pareto frontier, achieving 94.22% accuracy with 27.54 GB of peak training memory. The dashed line connects the Pareto-optimal methods. The vertical line marks the 80 GB H100 memory limit. Adam achieves 95.3% accuracy with 254.8 GB of peak memory. Right: Peak GPU memory decomposition of TACO on OPT-30B across 8 downstream tasks. The bars show model state, optimizer state, and other transient memory across all eight downstream tasks. TACO’s optimizer state occupies only 0.267 GB, less than 0.5% of peak memory on each task, while the transient memory accounts for the variation across tasks.
arXiv深度剖析
TACO 在维度归一化的 1→1 算子范数下给出精确最速下降方向：对二维权重矩阵的每一列，选取幅值最大元素的符号。 Muon 通过矩阵值更新降低优化器内存，但其几何与 AdamW 不同，微调 AdamW 预训练模型时可能性能下降；TACO 沿用 Muon 的算子范数最速下降视角，但把几何路线推进到列级一稀疏形式，同时保留一阶梯度。 摘要给出该方向的构造定义，并说明其保留一阶梯度、使优化器状态内存近乎可忽略；未提供推导细节或收敛性证明。
实用版 TACO 每列只维护少量低精度梯度分量，在 OPT-13B 上把持久优化器状态相对 AdamW8bit 减少 174 倍（27.7 GB 降至 0.16 GB），峰值训练内存减少 2.9 倍（80.6 GB 降至 27.5 GB），精度与运行时间相当。 既有方法或压缩优化器状态、或放弃一阶梯度、或在保留稠密状态的同时改变更新几何；TACO 在保留一阶梯度的前提下把状态压缩到近乎可忽略。 摘要报告了 OPT-13B 上的具体内存数值与“精度和运行时间相当”的对比结论，但未给出任务清单、评测指标或统计细节。
TACO 使 30-32B 参数模型能在单张 80 GB H100 上完成全参数微调，覆盖多个模型家族与任务。 此前全参数微调的优化器状态内存开销限制了可放入现代 GPU 的模型规模；TACO 把这一门槛推到 30-32B 量级。 摘要以“多个模型家族和任务”概括该结果，未列出具体模型名称、任务或评测分数。
启示与展望
该结果面向需要在有限 GPU 内存下做 LLM 全参数微调的研究者与工程师，适用场景是二维权重矩阵的优化器状态压缩；摘要所述证据来自 OPT-13B 的内存与精度对比，以及 30-32B 模型在单张 80 GB H100 上的多模型家族、多任务微调。若后续工作沿用其算子范数最速下降视角，可进一步探索列级一稀疏更新在其他优化几何与模型规模下的表现。
摘要未说明“精度相当”所依据的具体任务、指标与统计口径，也未给出 30-32B 模型微调所用的模型家族与任务清单；列级一稀疏更新在超出摘要所述设置时的行为、以及低精度梯度分量对训练稳定性的影响，仍是读者需要关注的问题。由于本次仅基于摘要，无法核对正文中的图表与实验细节。