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

summary

Video file (mp4)

The gist

The gist: TACO, a Ternary Absolute-max Column-wise One-sparse Optimizer, reduces optimizer state memory by 174× relative to AdamW8bit and achieves competitive accuracy for full-parameter fine-tuning

In short

TACO is a new first-order optimizer designed to drastically reduce memory overhead during full-parameter LLM fine-tuning. It achieves this by computing an exact steepest-descent direction using a ternary, column-wise, one-sparse approach based on the largest magnitude entries in weight matrices. This method reduces persistent optimizer state memory by 174 times compared to AdamW8bit while maintaining competitive accuracy.

Key concepts

TACO Optimizer
Ternary Absolute-max Column-wise One-sparse Optimizer is a first-order optimization method that computes updates based on the sign of the largest magnitude entries in columns of weight matrices. This geometric approach ensures updates are 'ternary, column-wise and one-sparse,' leading to nearly negligible persistent optimizer state memory.
Steepest Descent Direction
This is the direction in which a function decreases most rapidly at a given point. TACO computes this exact direction under a dimension-normalized 1->1 operator norm by selecting the sign of the largest magnitude entry in each column of two-dimensional weight matrices, providing an exact minimization path.
Operator Norm Geometry
This refers to a specific geometric constraint used to define the optimization problem. TACO reformulates local minimization as min D <G, D>F subject to ||D||(1,n->1,m) ≤ 1. This geometry dictates that the exact minimizer is defined by D⋆ = - m/n S(G), which defines the optimizer's update rule.
Memory Efficiency Gains
The practical design retains only a small set of low-precision gradient components per column, reducing persistent state memory complexity from O(mn) to O(kn). This results in a 174x reduction in memory compared to AdamW8bit and allows full-parameter fine-tuning of 30-32B models on single high-end GPUs.

Terminology used across episodes

This episode discusses

The paper

TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning · Read on arXiv

Jichao Jiang, Cristian McGee, El Houcine Bergou, HanQin Cai, Aritra Dutta

University of Central Florida · Mohammed VI Polytechnic University

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.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning".

Jane: The gist: TACO, a Ternary Absolute-max Column-wise One-sparse Optimizer,

Tom: First, who's behind it and why it matters.

Paper summary: Jane: So, looking at the whole picture of this "TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning" paper, what does it ultimately suggest about how we approach memory in training large language models?

Tom: It suggests that optimization geometry itself can be a framework for designing memory efficient first-order optimizers. Instead of just trying to compress the existing state or throwing away gradients, you can derive sparsity from the fundamental mathematical structure of the update problem.

Lu: I think what this points toward is that we need to look at how the local minimization problem is formulated, like they did with min D G, D F, D (1,n→1,m) at most one), and find a geometry that naturally yields a sparse solution <ref:2610.02199#pg1>.

Meng: For deployment at scale, this means we don't have to constantly sacrifice accuracy just to fit the model size on the hardware. We get this massive reduction in memory footprint without losing competitive performance on tasks like BoolQ or RTE, which is what we need for real-world training stability and cost management.

Jane: The implication is that as models get bigger and we are limited by GPU memory, designing optimizers based on these intrinsic geometric properties might be a more powerful way to manage the trade-off between memory efficiency and accuracy than simply tuning existing hyperparameters.

Conclusion: Tom: So we're wrapping up TACO, the Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning. The authors are working on making those large language models fit on GPUs by cutting down their optimizer memory a lot.

Jane: What's the main idea here, Tom? Like, what's the actual problem they’re trying to solve?

Tom: It’s that full fine-tuning of these big models eats up way too much memory for the optimizer state. They introduce this new optimizer called TACO which cuts that memory down by a factor of one hundred seventy-four compared to AdamW8bit.

Lu: The core mechanism they are using is based on finding the exact steepest-descent direction under a specific operator norm geometry, and they get sparsity from that structure itself.

Meng: From my side, when I think about deployment, what’s important is that this memory saving happens without hurting the accuracy of the fine-tuning run. That’s usually where things get shaky.

Jane: The paper claims they maintain comparable accuracy, so what are the actual numbers on how much memory it saves?

Tom: They show a reduction in peak training memory by nearly three times, from eighty point six gigabytes down to twenty seven point five gigabytes on an OPT thirteen billion parameter model.

Lu: And they prove a non-convex convergence guarantee for this vanilla TACO, showing the iterates actually reach an epsilon stationary point in order epsilon squared iterations.

Jane: That’s huge because proving convergence is often the hardest part of new optimization methods. So, what does this mean for someone who just wants to train a model without needing a PhD in math?

Tom: It means we can push full fine-tuning of models up to thirty two billion parameters on just one eighty gigabyte H100 GPU across different families and tasks.

Meng: That ability to do that on one card is what makes it viable for real-world production runs, not just research experiments.

Jane: It seems like the focus now shifts from finding a new optimizer every time to using the geometry of optimization itself as a tool to solve these resource constraints. Where does this leave us in terms of future work?

More episodes

← Home