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

TL;DR

SNIP is an adaptive mixed-precision training framework for large language models that optimizes layerwise precision to reduce computational cost while maintaining model quality, especially with subbyte precision support.

Contribution

It introduces a novel adaptive precision optimization method using metrics and ILP to improve training efficiency and stability for large language models.

Findings

Reduces FLOPs by up to 80% compared to baselines.

Maintains model quality across various model sizes and training phases.

Supports subbyte precision with minimal overhead.

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.