Gaussian Weight Sampling for Scalable, Efficient and Stable Pseudo-Quantization Training

TL;DR

The paper tackles the instability and computational burden of fully quantized training for large language models by proposing Gaussian weight sampling within pseudo-quantization training (PQT). It introduces a FP-friendly rounded Gaussian noise distribution and efficient noise generation to enable stable, scalable low-precision training across MX FP datatypes, achieving performance close to or surpassing BF16 in GPT2-124M and Llama2 models up to 1B parameters and 300B tokens with modest overhead. The key contributions include a formal Gaussian PQN framework, stochastic precision annealing, and practical implementation details that preserve forward–backward consistency while maintaining high throughput. This work lays a practical foundation for unified, low-precision quantization in large-scale model pre-training, offering significant efficiency gains with robust stability.

Abstract

Ever-growing scale of large language models (LLMs) is pushing for improved efficiency, favoring fully quantized training (FQT) over BF16. While FQT accelerates training, it faces consistency challenges and requires searching over an exponential number of cases, each needing over 200B tokens to ensure stability. Pseudo-quantization training (PQT) addresses the issues of FQT, although it is not well-studied. We explore the practical implications of PQT in detail and propose a noise distribution $R$ that is floating-point (FP)-friendly, with ideal properties including stochastic precision annealing. As a result, the proposed method serves as an effective theoretical foundation for low-precision FP parameters through PQT, utilizing efficient fake quantization via an addition and subsequent FP casting. We demonstrate that Gaussian weight sampling is (1) scalable: supports low-precision FP parameters down to FP6 and high-precision noise up to 9-bit with BF16 operator. The proposed method is (2) efficient: incurring computational overhead as low as 1.40\% on the A100 GPU in terms of Llama2 training tokens per second, and requiring 2 bytes per parameter in GPU memory. We demonstrate that PQT with Gaussian weight sampling is (3) stable: closely following or even surpassing performance of the BF16 baseline while pre-training GPT2 and Llama2 models with up to 1B parameters and 300B tokens.

Figures (8)

• Figure 1: Summary of Gaussian weight sampling (GaussWS).
• Figure 2: Example of effective PQN as in Equation \ref{['eq:gws_forward']} with $b_l=2$, $R=U(-0.5,0.5)$ and $b_t=4$ being underflow. Boxes wrapped in bold solid lines represent quantization groups with a square block size of 2, and an internal datatype of INT4 for $Q(\cdot)$. The second matrix represents PQN reflected in the backward pass while the fourth matrix represents the effective PQN during the forward pass.
• Figure 3: Training loss curve of the GPT2-124M model on the OpenWebText dataset.
• Figure 4: Training loss curve of the Llama2-{134M, 1B} models on the C4 dataset. First column represents average loss and the other two represent maximum loss. Third column corresponds to the range annotated with the orange arrow on the second column. For better visualization, weighted moving average is used with $\alpha=1/16$ on left column and $\alpha=1/128$ on right column.
• Figure 5: Resulting bitwidth $b_t$. Dots and red lines indicate layerwise mean and standard deviation. Upper and lower solid lines represent layerwise maximum and minimum while dotted lines represent transformer-blockwise maximum and minimum. Lines on 5, 9 and 12-bit divide the parameters into 3 tiers, and the percentages on the right-hand side represent the ratio of parameters for each. The order of layers is (qkv, out, up, down) for GPT2 and (q, k, v, out, gate, down, up) for Llama2.

Theorems & Definitions (8)

• Lemma 1
• proof
• Lemma 2
• Proposition 3
• Proposition 4
• proof
• proof
• proof