SNIP: An Adaptive Mixed Precision Framework for Subbyte Large Language Model Training

TL;DR

SNIP introduces a fine‑grained adaptive mixed‑precision framework for LLM pretraining that selects subbyte per‑layer formats (FP8/FP4) via an ILP optimization guided by two training‑quality metrics: forward loss divergence and backward weight divergence. By periodically collecting statistics, estimating divergences, and solving a knapsack‑like ILP (extended to handle pipeline parallelism), SNIP achieves substantial FLOP reduction (up to 80% FP4 usage) while maintaining near BF16 accuracy across 1B–70B models. The approach is implemented as an asynchronous system that integrates with standard training, incurring modest overhead and showing robust performance across checkpoints and model scales. Theoretical foundations link quantization perturbations to measurable losses and weight updates, enabling reliable per‑layer decisions and practical deployment in large‑scale LLM pretraining pipelines.

Abstract

Training large language models (LLMs) efficiently while preserving model quality poses significant challenges, particularly with subbyte precision supported by state-of-the-art GPUs. Current mixed-precision training approaches either apply uniform precision to all GEMM operations or rely on heuristic-based methods that fail to generalize during training, leading to suboptimal convergence and instability. To address these challenges, this paper introduces SNIP, a fine-grained adaptive mixed-precision training framework for LLM pretraining that supports subbyte precision. SNIP periodically collects statistics on activations, gradients, and optimizer states to assess the precision loss impact on model quality. We define two key metrics: loss divergence in the forward pass, caused by quantization-induced increases in training loss, and weight divergence in the backward pass, which measures error propagation through gradients affecting model updates. These metrics guide an Integer Linear Programming (ILP) problem that systematically optimizes layerwise precision to minimize overall quality loss while meeting efficiency targets. Experiments on 1B, 3B, 7B and 70B Llama-like models demonstrate that SNIP consistently outperforms existing baselines, reducing FLOPs by up to 80% while preserving model quality across different model sizes and training phases with minimal computational overhead.

Figures (13)

• Figure 1: Illustration of the gap of the training loss between high-precision (BF16) and low-precision (FPX). The gap consists of two parts: (1) Forward Loss Divergence: the increase in training loss directly introduced by quantization during the forward pass, and (2) Backward Weight Divergence: the accumulation of errors in weight updates during the backward pass due to quantization, which can compound over iterations and layers, impacting model convergence.
• Figure 2: System overview of SNIP, showing its integration into LLM training. SNIP periodically collects statistics on activations, weights, and optimizers, then asynchronously analyzes divergence metrics, solves an ILP problem, and updates layer-wise quantization.
• Figure 3: Comparison of accuracy versus efficiency (the fraction of FP4 FLOPs) for the TinyLlama 1B model. The FP8 baseline achieves the highest accuracy but lowest efficiency, while the FP4 configuration maximizes efficiency at the cost of accuracy. SNIP demonstrates a balance between high accuracy and a significant reduction in FLOPs by selectively applying FP4 to certain layers based on the loss and weight divergence metrics, outperforming other methods such as random and other heuristic-based approaches.
• Figure 4: The transformer block structure of Llama-like LLM. The blocks in blue (Q, K, V, O, Gate, Up and Down) are linear layers.
• Figure 5: The overall mixed precision training framework for linear layers, including the forward and backward pass.

Theorems & Definitions (4)

• theorem 1
• theorem 2
• definition 1
• definition 2