GAMMA: Global Bit Allocation for Mixed-Precision Models under Arbitrary Budgets

TL;DR

GAMMA is a post-training, quantizer-agnostic framework that efficiently allocates mixed-precision bits across modules of large language models, enabling high accuracy at reduced memory costs.

Contribution

It introduces a novel post-training method that learns module sensitivities and optimizes bit allocation via integer programming, avoiding costly retraining or static proxies.

Findings

GAMMA outperforms fixed-precision baselines by up to +12.99 in average accuracy.

It surpasses search-based mixed-precision methods by up to +7.00 in average accuracy.

GAMMA achieves fixed 3-bit quality at 2.5-bit average precision.

Abstract

Mixed-precision quantization improves the budget--accuracy trade-off for large language models (LLMs) by allocating more bits to sensitive modules. However, automating this allocation at LLM scale faces a unique combination of constraints: learnable approaches require quantization-aware training, which is infeasible for billion-parameter models; training-free alternatives rely on static proxy metrics that miss cross-module interactions and must be recomputed per target budget; and search-based methods are expensive without guaranteeing exact budget compliance. We propose GAMMA, a quantizer-agnostic framework that learns module-wise precision preferences entirely within a post-training pipeline. GAMMA optimizes a teacher-forced hidden-state reconstruction objective under an augmented Lagrangian constraint, and projects the learned preferences into exact budget-feasible discrete assignments via integer programming. A key property is score reuse: because the learned preferences encode a stable sensitivity ranking rather than budget-specific weights, a single training run serves arbitrary deployment targets by re-solving only the integer program, reducing per-budget adaptation from hours to a few minutes. Across Llama and Qwen models (8B--32B), GAMMA outperforms both fixed-precision baselines (up to +12.99 Avg.) and search-based mixed-precision methods (up to +7.00 Avg.), and can match fixed 3-bit quality at 2.5-bit average precision, enabling deployment at substantially smaller memory footprints.