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

TL;DR

This paper introduces Gaussian weight sampling for pseudo-quantization training, enabling scalable, efficient, and stable low-precision training of large language models with minimal computational overhead.

Contribution

It proposes a novel Gaussian weight sampling method for pseudo-quantization training that improves scalability, efficiency, and stability in training large language models.

Findings

Supports low-precision FP parameters down to FP6

Incurred only 1.40% overhead on A100 GPU

Achieved stable training surpassing BF16 baseline on GPT2 and Llama2

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.