TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning

2026-10-01

AI Summary

This study addresses the excessive memory overhead of optimizer states during full-parameter fine-tuning of large language models by proposing TACO, a novel optimizer grounded in the steepest descent perspective under operator norms. TACO employs a ternary maximum-absolute-value column-sparse algorithm combined with low-precision gradient storage to preserve first-order gradient information while achieving near-zero persistent states, thereby overcoming traditional geometric constraints to balance accuracy and efficiency. Experimental results demonstrate that TACO reduces peak memory consumption by 2.9× and compresses state memory by 174× when training OPT-13B. Furthermore, it enables efficient fine-tuning of 30–32B parameter models on a single H100 GPU without compromising model performance.

Abstract

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. Still, its geometry differs from AdamW and can lead to performance degradation when fine-tuning AdamW-pretrained models. To reduce optimizer memory without sacrificing accuracy or computational efficiency in LLM fine-tuning, we propose Ternary Absolute-max Column-wise One-sparse optimizer, or TACO, which follows Muon's operator-norm steepest-descent view but takes the geometric route further. TACO computes the exact steepest-descent direction under a dimension-normalized $1\to1$ operator norm by selecting the sign of the largest magnitude entry in each column of two-dimensional weight matrices. This retains first-order gradients while making optimizer state memory nearly negligible. Our practical TACO optimizer maintains only a small set of low precision gradient components per column, reducing persistent optimizer state by $174\times$ relative to AdamW8bit (from 27.7 GB to 0.16 GB) and peak training memory by $2.9\times$ (from 80.6 GB to 27.5 GB) on OPT-13B, while achieving comparable accuracy and runtime. TACO further enables full-parameter fine-tuning of 30-32B-parameter models on a single 80 GB H100 GPU across multiple model families and tasks.

Problem

• LLM fine-tuning
• optimizer memory overhead
• full-parameter fine-tuning
• optimizer state compression

Innovation

• Optimizer State Compression
• Operator-norm Steepest Descent
• Full-parameter Fine-tuning
• Memory Efficiency
• Large Language Models