TACO: Ternary Absolute-max Column-wise One-sparse Optimizer Cuts LLM Fine-Tuning Optimizer Memory by 174x

Forum topic · 小凯 · 2026-10-04

Summary

TACO (Ternary Absolute-max Column-wise One-sparse optimizer) is a new optimizer for full-parameter fine-tuning of large language models that dramatically reduces optimizer state memory. Building on Muon's operator-norm steepest-descent view, TACO computes the exact steepest descent direction under a dimension-normalized 1→1 operator norm: for each column of a 2D weight matrix it keeps only the sign of the largest absolute element. This preserves first-order gradient information while making optimizer state memory nearly negligible. In practice, TACO maintains only a few low-precision gradient components per column. On OPT-13B, it reduces persistent optimizer state 174x compared to AdamW8bit (27.7 GB to 0.16 GB) and lowers peak training memory 2.9x (80.6 GB to 27.5 GB), while maintaining comparable accuracy and runtime speed. TACO also enables full-parameter fine-tuning of 30-32B parameter models on a single 80 GB H100 GPU across multiple model families and tasks. The paper by Jiang et al. is available as arXiv:2610.02199.

论文概要

研究领域: ML 作者: Jichao Jiang, Cristian McGee, El Houcine Bergou, Hanqin Cai, Aritra Dutta 发布时间: 2026-10-01 arXiv: 2610.02199

Overview

Full-parameter fine-tuning of large language models (LLMs) incurs substantial optimizer state memory overhead, limiting the model sizes that fit on modern GPUs. Existing approaches either compress optimizer state, abandon first-order gradients, or change the update geometry while retaining dense state. The recently introduced Muon optimizer reduces optimizer memory through matrix-valued updates, but its geometry differs from AdamW and can degrade performance when fine-tuning AdamW-pretrained models.

The TACO Approach

To reduce optimizer memory without sacrificing accuracy or computational efficiency in LLM fine-tuning, the authors propose the Ternary Absolute-max Column-wise One-sparse optimizer (TACO). TACO follows Muon's operator-norm steepest-descent view but takes the geometric route further: it computes the exact steepest descent direction under a dimension-normalized 1→1 operator norm — the sign of the absolute-maximum element of each column of a 2D weight matrix.

Key properties:

Retains first-order gradient information
Makes optimizer state memory nearly negligible
Each column only needs a few low-precision gradient components in the practical optimizer

Results

On OPT-13B, versus AdamW8bit: persistent optimizer state reduced 174x (27.7 GB → 0.16 GB)
Peak training memory reduced 2.9x (80.6 GB → 27.5 GB)
Maintains comparable accuracy and running speed
Enables full-parameter fine-tuning of 30-32B parameter models on a single 80 GB H100 GPU, across multiple model families and tasks