TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning
TACO attacks the memory bottleneck of full-parameter LLM fine-tuning by replacing dense AdamW states with a geometry-induced sparse update. The core idea is to use a dimension-normalized $\ell_1 \to \ell_1$ operator norm, whose exact steepest-descent direction is a ternary, column-wise one-sparse matrix. By retaining only a tiny FP8 heavy-hitter history per column, TACO slashes persistent optimizer state by over $170\times$ and peak memory by nearly $3\times$, enabling 30–32B parameter models to fine-tune on a single 80 GB GPU.
TACO is a conceptually fresh optimizer that successfully trades off memory against dense gradient information. The geometric derivation is clean, the empirical memory savings are large, and the multi-architecture evaluation is convincing. However, the elegant theory largely describes a ``vanilla'' history-free update that differs noticeably from the practical FP8 heavy-hitter implementation, and several strong assumptions underlie the claims about Adam-like continuation geometry and $\mu$P transfer. It is a strong systems-minded paper with solid theory, though readers should view the theoretical guarantees as stylized characterizations rather than airtight justifications for the deployed algorithm.
The derivation of the exact steepest-descent direction under the dimension-normalized $\| \cdot \|_{(1,n \to 1,m)}$ operator norm is rigorous and elegant; it naturally yields an update with at most one nonzero per column without post-hoc sparsification. The memory engineering is equally compelling: the sparse FP8 gradient history reduces state complexity from $\mathcal{O}(mn)$ to $\mathcal{O}(kn)$ with $k=16$, and the experiments demonstrate that this design scales to 32B-parameter models on a single H100 while preserving competitive downstream accuracy. The theoretical result that TACO and an Adam proxy share the same continuation path in a fixed-input setting offers a principled reason to expect less optimizer mismatch than Muon.
The gap between theory and practice is notable. The convergence guarantee applies to ``vanilla'' history-free TACO, whereas the deployed algorithm relies on an EMA of sparse heavy hitters ($\mu=0.95$, $k=16$) and auxiliary AdamW for non-matrix parameters—engineering choices that are only loosely connected to the geometric derivation. The continuation geometry and $\mu$P analyses rely on strong, unverified assumptions: Theorem 2 assumes a fixed-input coherent objective and common target output, while Theorem 4 requires ``balanced winner counts'' ($\max_i c_i = \Theta(n/m)$) that the paper does not empirically validate on real gradients. Additionally, the default 2,000-step budget undershoots TACO's best performance (BoolQ and RTE improve by several points at 16K steps), so the main comparisons may understate its true cost in wall-clock time. Finally, the throughput is only middling: FlashAdamW trains roughly twice as fast in tokens per second, raising the question of whether TACO's memory savings fully compensate for its iteration-level overhead when training to convergence.
The empirical evidence robustly supports the memory claims, with detailed decompositions showing TACO's persistent state falls below 0.5% of peak memory. Accuracy is broadly competitive—on OPT-13B SST-2, TACO reaches 94.2% versus 95.3% for AdamW—though it is not uniformly Pareto-optimal when wall-clock throughput is considered (FlashAdamW achieves 547 tokens/s versus TACO's 276). The comparisons are reasonably fair in that each optimizer uses its own tuned learning rate and standard configurations, but the lack of a unified hyperparameter sweep or training-to-convergence protocol makes it hard to separate optimizer quality from tuning effort. The paper wisely focuses on first-order memory-efficient baselines rather than zeroth-order methods, which avoids conflating memory savings with the cost of extra forward passes.
Reproducibility is mixed. The authors release source code and provide detailed hyperparameters in Appendix B, including batch size, learning rate schedules, and FP8 formats. However, the implementation relies on PyTorch autograd hooks and custom in-place sparse updates that are hardware- and software-version sensitive, and the paper notes that FP8 E4M3 and int32 indices are used without specifying how NaNs or rounding modes are handled. The ``Anonymous source code'' phrasing in the reproducibility statement suggests the repository may be provisional or de-identified for review. Exact reproduction would require access to the same 80 GB H100 nodes and careful matching of the low-level gradient checkpointing and hook implementation.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.