Q-PACE: Dynamic Precision Allocation for Quantization-Aware Training

The gist: Q-PACE proposes a new approach, called Q-PACE, consisting of a second-order sensitivity model that predicts the loss increase as a sum of quantization noise MSE weighted by per-layer curvature coefficients, which consistently improves over existing mixed-precision training recipes and achieves comparable loss at substantially lower total memory budgets

The Problem: QAT with Adaptive Mixed-Precision

The goal is to choose the quantization levels of n layers in a network during quantization-aware training The problem formally seeks to find B∗ = argmin B=(b1,..., bn) L(W, B) subject to X i bidi ≤ C0 (Equation 1), where L(W, B) is the loss and C(W, B) is the resource cost For QAT, we are interested in finding a sequence of allocations B1,..., BT that ultimately achieves the lowest final loss while adhering to the cost budget at each iteration

Methodology: Q-PACE Procedure

Q-PACE involves estimating per-layer sensitivity coefficients αi using only forward passes and then solving a knapsack problem to determine “optimal” per-layer bit-widths

The procedure outlined in Algorithm 1 requires evaluating Lhigh on D, adding noise δWi ∼ N(0, σ2k) to layer weight Wi, determining MSE([Wi, Wi + δWi), and estimating sensitivity αi via least squares using Equation (6)

The expected loss change is approximated as E[L(W + δW) − L(W)] ≈ X i αi MSE([Q(Wi, Bi), Wi), (5), where Lhigh = L(W) is the loss over full-precision weights and αi is an estimate of 1/2 Tr(Hii) that can be interpreted as the aggregated curvature of a layer

The bitwidth assignment is then solved by solving Equation (1) using Equation (5) as a proxy for the intractable L(W, B), which constitutes a multiple choice knapsack problem (MCKP)

Key Findings on Performance and Trade-offs

Q-PACE consistently improves over existing mixed-precision training strategies across model scales, precision budgets, and format options In pretraining, it achieves the lowest final loss in every setting evaluated, recovering up to 81% of the loss gap between uniform NVFP4 and MXFP8 training

Q-PACE provides a substantially better quality–precision trade-off, producing better final loss than both SNIP and the NVIDIA heuristic at a fixed precision budget, while at matched loss it requires fewer bits per model parameter

The quantization sensitivity is highly structured across the network, depending on both layer type and depth, and its profile evolves over the course of training

Sensitivity Analysis and Dynamics

Sensitivity is structured by layer type and depth, with Q-PACE promoting five of the six quantized linear layers to full precision in the final block Sensitivity evolves during training; the largest changes occur between the early and middle stages of training, after which the allocation becomes substantially more stable

The relative sensitivity by layer type over training for the 800M model shows that Q-PACE assigns the largest relative sensitivity to the MLP projections, whereas SNIP places comparatively greater sensitivity on K projections Sensitivity is also strongly structured by model depth, with later blocks generally being more sensitive for Q-PACE

Computational Overhead and Practicality

Q-PACE calibration requires only forward passes, resulting in negligible memory overhead compared to gradient-based methods One recalibration costs mn + 1 forward passes over D, and the relative compute overhead is linear in both m and D and inversely proportional to τ In experiments, the measured wall-clock time overhead was only 5% for a 4B model corresponding to a conservative choice of τ = 5k

Conclusion

Q-PACE is a new method for adaptive mixed-precision LLM training that probes linear layers and estimates their sensitivity, which is then used for periodic reassignment of precision across layers during training With no memory overhead, our method consistently improves over existing baselines and allows for substantial memory savings at a matched performance Across pretraining experiments from 350M to 4B parameters and supervised fine-tuning of Qwen3-4B, it outperforms other methods and recovers up to 81% of the NVFP4-to-MXFP8 loss gap at 6.5 bit/parameter

How it works

The core mechanism relies on modeling the quantization error as a sum of per-layer errors weighted by sensitivity coefficients This is formalized in Equation (5), where L(W, B) ≈ Lhigh + X i αi MSE([Q(Wi, Bi), Wi), where Lhigh = L(W) is the loss over full-precision weights and αi is an estimate of 1/2 Tr(Hii) that can be interpreted as the aggregated curvature of a layer

The sensitivity coefficients αi are estimated via a forward-pass calibration procedure, which is viewed as a variant of Hutchinson trace estimation, where the layerwise curvature αi is estimated from the increase in loss induced by isotropic random perturbations

Bitwidth Assignment

Once sensitivities αi are determined, the optimal sequence of allocations (B1,..., BT) is found by solving Equation (2), which involves minimizing L(WT(B1,..., BT-1), BT) subject to max t=1,… T C(Bt) ≤ C0 This discrete selection problem is solved using pseudo-polynomial dynamic programming (Pisinger, 1995), where the budget C0 is discretized into M bins and each layer cost dibi is rounded up to the whole number of bins

MSE Estimation

The estimate MSE([can be obtained empirically by measuring the statistics of the actual weight tensors’ quantization error, but we rely on the efficient analytic approximation, MSE([Q(Wi), Wi) ∝ 4-bi, based on the following analysis This analytic approximation recovers expressions of high-resolution quantization (Gish & Pierce, 1968)

Alternative Sensitivity Models

The paper compares Q-PACE against several alternative approximations, including Q-PACE inversed, which uses a symmetric linear approximation relative to the loss of uniformly FP4 quantized model Other models include the first-order Taylor sensitivity (Equation 27) and cosine similarity (Equation 28), showing that Q-PACE remains the strongest variant in Table 7

Downstream Task Scores

Q-PACE achieves superior performance on downstream benchmarks, such as GSM8K, IFEval, HumanEval, Natural Questions (NQ), and TruthfulQA across all tested budgets For instance, at budget 6.5 bits/parameter for the 4B model with menu set to 4 and 16, Q-PACE achieves an accuracy of 80.1±0.

1. Bold header: Adaptive Layer Precision Allocation during Training

The Q-PACE method consistently improves over existing mixed-precision training strategies across model scales, precision budgets, and format options, allowing it to recover up to 81% of the loss gap between uniform NVFP4 and MXFP8 training. This enables the system to achieve substantially lower total memory budgets while maintaining high accuracy during pretraining.

2. Bold header: Dynamic Sensitivity Recalibration

Q-PACE re-estimates layer sensitivities periodically during training, and because sensitivity evolves during training, it proposes that only a small number of reallocations (on the order of 10 over a full training run) is sufficient, keeping the overhead of Q-PACE low in practice. This allows for dynamic adaptation to changing network dynamics without excessive computational cost.

3. Bold header: Optimized Hardware Utilization

The system can be configured to utilize specific hardware features by testing formats like NVFP4, MXFP6, MXFP8, and the results show that Q-PACE outperforms baselines across these options at a fixed precision budget. This allows for optimal mapping of model operations to available low-bit data formats on supported accelerators.

4. Bold header: Robust Performance Across Formulations

The method is robust across alternative sensitivity models, as shown by Table 7, where the default Q-PACE formulation remains the strongest variant, suggesting that the core mechanism—a second-order loss model as a sum of per-layer errors weighted by sensitivity coefficients—is fundamentally sound regardless of whether it uses forward error only or other approximations.

5. Bold header: Inference Memory Trade-off Optimization

Q-PACE forms an empirical Pareto frontier, meaning it can achieve lower loss, or equivalently reach a comparable loss with less memory across various inference formats like WmixAmix and W4A8. This allows the system to select the optimal precision allocation based on the required inference memory budget.