Compute-Optimal Quantization-Aware Training (2509.22935v1)

Abstract: Quantization-aware training (QAT) is a leading technique for improving the accuracy of quantized neural networks. Previous work has shown that decomposing training into a full-precision (FP) phase followed by a QAT phase yields superior accuracy compared to QAT alone. However, the optimal allocation of compute between the FP and QAT phases remains unclear. We conduct extensive experiments with various compute budgets, QAT bit widths, and model sizes from 86.0M to 2.2B to investigate how different QAT durations impact final performance. We demonstrate that, contrary to previous findings, the loss-optimal ratio of QAT to FP training increases with the total amount of compute. Moreover, the optimal fraction can be accurately predicted for a wide range of model sizes and quantization widths using the tokens-per-parameter-byte statistic. From experimental data, we derive a loss scaling law that predicts both optimal QAT ratios and final model performance across different QAT/FP compute allocation strategies and QAT bit widths. We use the scaling law to make further predictions, which we verify experimentally, including which QAT bit width is optimal under a given memory constraint and how QAT accuracy with different bit widths compares to full-precision model accuracy. Additionally, we propose a novel cooldown and QAT fusion approach that performs learning rate decay jointly with quantization-aware training, eliminating redundant full-precision model updates and achieving significant compute savings. These findings provide practical insights into efficient QAT planning and enable the training of higher-quality quantized models with the same compute budget.

• The paper introduces a compute-optimal QAT framework that derives the optimal quantization phase fraction using a log-linear function of the tokens-per-parameter-byte statistic.
• The unified loss scaling law models loss as a function of model size, token counts, and bit-width, enabling precise prediction of performance and compute waste.
• The proposed cooldown & QAT fusion technique streamlines training by jointly decaying the learning rate with quantization, yielding significant compute savings.

Compute-Optimal Quantization-Aware Training: Theory, Empirics, and Practical Implications

Introduction

This paper presents a comprehensive analysis of compute allocation strategies for quantization-aware training (QAT) in LLMs, challenging prevailing assumptions about the optimal division of training between full-precision (FP) and quantized phases. The authors provide a rigorous empirical and theoretical framework for determining the compute-optimal QAT fraction, introduce a unified loss scaling law for QAT, and propose a novel cooldown & QAT fusion technique for improved training efficiency. The work is grounded in extensive experiments across a wide range of model sizes (86M–2.2B parameters), compute budgets, and quantization bit-widths (1–6 bits), with strong generalization to different datasets and hyperparameters.

Problem Formulation and Motivation

QAT is the dominant approach for adapting neural networks to low-precision inference, outperforming post-training quantization (PTQ) in accuracy, especially at low bit-widths. The standard practice is to pretrain a model in FP, then fine-tune with QAT. However, the optimal allocation of compute between FP and QAT phases has been poorly understood, with prior work suggesting a fixed QAT fraction (e.g., 10%) regardless of model size or compute budget. This paper demonstrates that such fixed ratios are suboptimal, especially as compute budgets increase, and that the optimal QAT fraction is a function of the tokens-per-parameter-byte statistic, which encapsulates model size, training duration, and quantization width.

Empirical Findings: Scaling Laws and Optimal QAT Fraction

The authors conduct a systematic grid of experiments, varying FP/QAT token allocation, model size, and bit-width. Key findings include:

• The optimal QAT fraction increases with total compute budget (tokens-per-parameter-byte), contradicting previous claims of a fixed optimal ratio.
• Low-bit QAT (1–2 bits) is highly sensitive to the QAT fraction, with substantial compute waste if suboptimal allocation is used. For 1-bit QAT, up to 50% compute savings are possible by using the optimal fraction.
• Larger models tolerate lower bit-widths for QAT and can match FP accuracy at lower precision, given sufficient compute.
• The optimal QAT fraction decreases with increasing model size for fixed compute, but increases with decreasing bit-width.

The optimal QAT fraction $f^*$ is accurately predicted by a log-linear function of the tokens-per-parameter-byte statistic:

$$f^*(D_{total}, N, B) = \exp\left(a \cdot \log\left(\frac{D_{total}}{N \cdot B}\right)\right)$$

where $a$ is a fitted parameter.

Unified Loss Scaling Law

A central contribution is the derivation of a unified loss scaling law for QAT, modeling final expected loss as a function of model size ($N$), FP/QAT token counts ($D_{fp}, D_{qat}$), and bit-width ($B$):

$$L(N, D_{qat}, D_{fp}, B) = \text{Chinchilla-like loss} + \text{QAT fraction-aware penalty}$$

The penalty term incorporates irreducible QAT error, pure QAT penalty, and FP/QAT interaction, all parameterized by tokens-per-parameter-byte. The law achieves high fit quality ($R^2 = 0.982$–$0.991$ across bit-widths) and enables:

• Prediction of optimal QAT fraction for arbitrary compute budgets and model sizes
• Quantification of compute waste for suboptimal QAT allocation
• Determination of the minimal bit-width required to match FP accuracy for a given model and compute budget
• Guidance for parameter-precision trade-offs under memory constraints

Cooldown & QAT Fusion: Training Efficiency

The paper introduces a cooldown & QAT fusion technique, where learning rate decay is performed jointly with QAT rather than sequentially after FP training. This approach eliminates redundant FP updates that are effectively discarded during QAT initialization, yielding:

• Consistent improvements in perplexity for 4- and 6-bit QAT across all model sizes and token counts
• Compute savings quantified in "wasted tokens" units, with up to 13% savings for large models
• Marginal improvements for 1- and 2-bit QAT, attributed to the large optimal QAT fraction in these regimes

Implementation Details

Model and Training Setup

• Architecture: Decoder-only transformer (Llama 2-like), SwiGLU, RoPE, RMSNorm, alternating attention/FFN, tied embeddings/LM head.
• Optimizer: Adam with decoupled weight decay, bfloat16 mixed precision.
• Dataset: DCLM, tokenized with Llama 2 tokenizer, concat-and-chunk batching.
• QAT Algorithms: ParetoQ (state-of-the-art), with Elastic Binarization (1-bit), SEQ (2-bit), LSQ (≥3-bit), per-output-feature quantization scales.
• Learning Rate Scheduling: Warmup-stable-decay (WSD) for FP, cosine decay for QAT, with fusion as described.

Practical Guidelines

• Compute Allocation: Use the scaling law to determine $f^*$ for given $N$, $B$, and $D_{total}$.
• Bit-Width Selection: For a fixed memory budget, select the lowest bit-width that matches FP accuracy, as higher bit-widths do not improve accuracy but increase memory usage.
• Training Efficiency: Employ cooldown & QAT fusion to maximize compute utilization, especially for mid/high bit-width QAT.
• Generalization: The scaling law and optimal fraction prediction generalize to larger models (2.2B) and different datasets (SlimPajama), with minimal sensitivity to hyperparameters.

Theoretical and Practical Implications

Theoretical

• Scaling laws for QAT are not static; they depend on compute, model size, and bit-width.
• Unified loss scaling law enables principled resource allocation and model selection for quantized LLMs.
• Parameter-precision trade-off analysis provides a framework for memory-constrained deployment, relevant for on-device and edge inference.

Practical

• Suboptimal QAT allocation can result in substantial compute waste, especially for low-bit QAT.
• Cooldown & QAT fusion should be adopted as standard practice for efficient QAT pipelines.
• Practitioners can use the provided scaling law to plan training runs, select bit-widths, and optimize for deployment constraints.

Future Directions

• Interaction of pretraining precision (e.g., FP8, FP4) with QAT scaling laws
• Extension to multi-stage training pipelines (SFT, RL, multimodal) and their impact on QAT allocation
• Generalization to other architectures and modalities

Conclusion

This work provides a rigorous framework for compute-optimal QAT in LLMs, overturning previous assumptions about fixed QAT allocation ratios and introducing a unified loss scaling law for principled resource planning. The empirical and theoretical results have direct implications for efficient training and deployment of quantized models, especially in memory- and compute-constrained environments. The proposed cooldown & QAT fusion technique further enhances training efficiency. The findings are robust across model sizes, datasets, and hyperparameters, and the methodology is extensible to future advances in low-precision training and multi-stage pipelines.