Automatic mixed precision for optimizing gained time with constrained loss mean-squared-error based on model partition to sequential sub-graphs

TL;DR

The work tackles automatic mixed-precision configuration for post-training quantization under a constrained loss, introducing a loss MSE metric derived from a first-order Taylor expansion of quantization noise and showing it decomposes additively across layers. It also introduces a hardware-aware time-gain model by partitioning the model into sequential sub-graphs, measuring empirical time gains per group, and optimizing via Integer Programming to maximize time (or memory) gain within a loss-MSE threshold. The method is validated on Intel Gaudi 2 accelerators across several LLMs, comparing IP-based strategies (IP-ET, IP-TT, IP-M) against Random and Prefix baselines, and consistently achieving better accuracy-latency and accuracy-memory tradeoffs. This yields a practical, scalable approach to MP configuration for PTQ with strong empirical support on real hardware and diverse language tasks. The combination of a Taylor-based sensitivity metric, group-wise additivity of time gains, and IP optimization offers a robust framework for hardware-aware quantization in edge and data-center inference scenarios.

Abstract

Quantization is essential for Neural Network (NN) compression, reducing model size and computational demands by using lower bit-width data types, though aggressive reduction often hampers accuracy. Mixed Precision (MP) mitigates this tradeoff by varying the numerical precision across network layers. This study focuses on automatically selecting an optimal MP configuration within Post-Training Quantization (PTQ) for inference. The first key contribution is a novel sensitivity metric derived from a first-order Taylor series expansion of the loss function as a function of quantization errors in weights and activations. This metric, based on the Mean Square Error (MSE) of the loss, is efficiently calculated per layer using high-precision forward and backward passes over a small calibration dataset. The metric is additive across layers, with low calibration memory overhead as weight optimization is unnecessary. The second contribution is an accurate hardware-aware method for predicting MP time gain by modeling it as additive for sequential sub-graphs. An algorithm partitions the model graph into sequential subgraphs, measuring time gain for each configuration using a few samples. After calibrating per-layer sensitivity and time gain, an Integer Programming (IP) problem is formulated to maximize time gain while keeping loss MSE below a set threshold. Memory gain and theoretical time gain based on Multiply and Accumulate (MAC) operations are also considered. Rigorous experiments on the Intel Gaudi 2 accelerator validate the approach on several Large Language Models (LLMs).

Figures (9)

• Figure 1: Measured empirical time gain $c^{\textrm{ET}}_{j,p}$ of the Attention sub-graph in Llama-3.1-8B (in blue) compared to its prediction based on the summation of per-layer time gain measurements (in orange) and for the theoretical time gain $c^{\textrm{TT}}_{j,p}$ (green) for any of its $2^{5}$MP configurations. The various configurations are ordered in ascending order of empirical time gain. Configurations are labeled as 5-bit binary words which represent the numerical format of each of the 5 linear operations ( q_proj, v_proj, k_proj, qk_matmul, and av_matmul) with BF16 and FP8 denoted as 0 and 1, respectively.
• Figure 2: Layer-wise quantization patterns across MP configurations (rows) and model layers (columns) for IP-ET (top), Prefix (middle), and Random (bottom). Yellow: FP8, purple: BF16.
• Figure 3: Empirical validation of the additivity assumption on different MP configurations
• Figure 4: Theoretical loss MSE vs. empirical time gain on the 1B model across four tasks.
• Figure 5: Average accuracy difference [%] vs. TTFT across HellaSwag, LAMBADA, Winogrande, and PIQA. Comparing MP quantization strategies (IP-ET, Random and Prefix)