Exploring Quantization for Efficient Pre-Training of Transformer Language Models (2024)

Abstract: The increasing scale of Transformer models has led to an increase in their pre-training computational requirements. While quantization has proven to be effective after pre-training and during fine-tuning, applying quantization in Transformers during pre-training has remained largely unexplored at scale for language modeling. This study aims to explore the impact of quantization for efficient pre-training of Transformers, with a focus on linear layer components. By systematically applying straightforward linear quantization to weights, activations, gradients, and optimizer states, we assess its effects on model efficiency, stability, and performance during training. By offering a comprehensive recipe of effective quantization strategies to be applied during the pre-training of Transformers, we promote high training efficiency from scratch while retaining language modeling ability.

## introduction transformers @xcite have become the dominant model for natural language processing, with the gpt family of models @xcite showcasing their effectiveness across various tasks, from language understanding to code generation. as the performance of transformers scales nicely with the number of parameters and the data size, the current state-of-the-art models have reached unprecedented computational requirements both during training and inference @xcite . for instance, pre-training a 175b parameter gpt-3 model requires a staggering number of 10,000 v100 gpus for 14.8 days @xcite . as a result, pre-training has become extremely expensive, beyond the reach of most research groups, and has raised concerns over sustainability due to co2 emissions from extensive gpu usage @xcite . to improve the efficiency of transformers, quantization has gained significant traction due to its recent successes in both post-training @xcite @xcite and during fine-tuning @xcite . quantized pre-training, where certain parts of the computational graph and model parameters are quantized from the beginning of training, remains a challenging problem. in such scenarios, the training instabilities caused by substantial changes in model parameters and emerging model behaviors do not pair well with the added noise introduced by quantization @xcite . additionally, quantizing model components without compromising performance becomes increasingly difficult at larger scales @xcite . despite its importance, quantized pre-training of transformer language models remains largely unexplored at scale. in this paper, we present the first in-depth study on the effects of quantizing transformer language models during pre-training and at scale. our primary aim is to provide a recipe for quantized pretraining by conducting a controlled study that investigates the impact of quantization on weights, activations, gradients, and optimizer states on model efficiency, stability, and performance, using a simple linear quantization with 4 and 8 bits. our findings demonstrate that 8-bit quantization for weights and activations can be effectively combined to provide significant memory savings and potential speedup, achieving performance comparable to the baseline model. however, extending quantization to gradients to utilize computational speedup in backward matrix multiplications or reducing precision to 4 bits results in notable training instability. specifically, 4-bit quantization introduces a sharper loss landscape for weights ( §4.1) and persistent outliers in the channel dimension of activations ( §4.2), significantly degrading performance despite attempts to manage them through per-channel quantization. additionally, gradient quantization is particularly problematic due to spikes in gradient norms during early training phases and the unstructured and sparse nature of gradients throughout training ( §4.3). while 8-bit gradient quantization does not hurt model convergence, transitioning to 4 bits results in non-convergence. additionally, quantizing the first-order moments of adam to even 4 bits is feasible without significant performance loss, but the second-order moments require a more complex quantization scheme to avoid instabilities in the adam update, even when using 8-bits quantization ( §4.4). lastly, we present our recommended pre-training quantization recipe for the different model components ( §4.5). ## related work in recent years, numerous methods have been studied to improve the efficiency of neural networks. among these, quantization-aware training (qat) emerges as an acceleration technique for inference, as parameters are stored and operations are conducted with higher precision during training. the induced quantization error during training serves as a regularizer, as demonstrated in @xcite , ultimately facilitating the development of a more quantization-friendly model. in contrast, fully quantized training (fqt) harnesses the accelerated gains derived from higher throughput in int8 or int4 operations supported by modern gpus during training. additionally, fqt capitalizes on memory savings by storing parameters in lower precisions. as exemplified by @xcite and @xcite , states of the adam optimizer are stored in 4 and 8 bits, respectively, to minimize memory footprint. another strategy introduced by @xcite involves quantizing both weights and gradients to reduce bandwidth usage in distributed training. furthermore, @xcite and @xcite advocate the replacement of linear operations with int8 matrix multiplications to achieve substantial speedups. our work extends quantization to multiple components, providing a more comprehensive exploration of the challenges and benefits of quantization during pre-training. despite the considerable advantages in accelerating the training process, fqt poses challenges attributable to numerical stability and optimization issues inherent in training quantized networks. many existing methods predominantly focus on fine-tuning large language models (llms) using fqt, leveraging the inherent stability of pre-trained models in each gradient update. the evolution from 16-bit to fp8 data formats, as evidenced by remarkable results in mixed precision training on llms @xcite , showcases the potential of fqt. however, the scarcity and difficulty in obtaining hardware that supports fp8 formats pose significant challenges to its widespread adoption. @xcite employ 8-bit quantization for linear operations in both forward and backward passes, achieving remarkable results in pretraining large-scale vision language models. their approach incorporates row-wise quantization for activations and gradients, mitigating the impact of quantization errors on other parameters. however, vision language models significantly differ from large textual language models. @xcite explored 4-bit quantization in fully quantized training (fqt) using hadamard transformations to handle activation outliers and proposed bit splitting to quantize gradients in 4-bit precision. however, their work primarily focuses on fine-tuning, and only relatively small models (60m parameters) are pre-trained on a small dataset (wmt 14 en-de), which hinders the generalizability of their findings to larger language models. our study explores pretraining at a larger scale, utilizing a bigger model and dataset to thoroughly examine the effects of quantization on various model components. the generalization of these findings to largescale language models (llms) remains a challenge, especially considering that training such models involves unique complexities. additionally, their reliance on hadamard transformations imposes restrictions on activation dimensions, limiting applicability to power-of-two dimensions, a constraint not easily met by recent llm architectures such as llama. in contrast to the aforementioned complex quantization methods, such as quantile quantization and learnable quantization parameters, this study focuses on a more straightforward implementation. the objective is to investigate the effects of quantization and explore the feasibility of training llms with full integer operations. the primary goal is to offer detailed insights into quantizing different components of the model. while more sophisticated quantization methods may enhance performance, our work serves as a foundational investigation, providing insights and paving the way for further exploration in quantization for llms. ## quantization methodology in this work, we explore various quantization schemes ( §3.1) and granularities ( §3.2) on model components. this controlled approach during pretraining allows us to examine the impacts on language modeling and downstream task performance. ## experimental results we used gpt-2 small (124m) @xcite with flashattention-2 @xcite for our experiments due to its popularity as a baseline architecture for contemporary studies of llms. while larger models are often used, gpt-2 small provides a manageable framework for in-depth experimentation without requiring the immense computational resources of larger models. for our experiments, we pre-trained 30 models from scratch on open-webtext @xcite for 300k gradient steps with a global batch size of 512 samples and a context length of 1024 tokens, processing approximately 157 billion tokens, which is consistent with similar works @xcite . for our evaluation setup, we evaluate the performance of the models on a range of language tasks, including arc-easy @xcite , arc-challenge @xcite , hellaswag @xcite , lambada @xcite , and glue @xcite . additional details about our training and evaluation setups are provided in appendix a. ## conclusion this study presents an extensive analysis of the impact of quantizing specific transformer components in 4 and 8 bits during pre-training, in contrast to concurrent work that focuses on individual methods without comprehensive ablations @xcite . our study reveals that quantizing weights to 8 bits from the beginning of pre-training is generally successful. however, 4-bit weight quantization can significantly affect model convergence due to a sharper loss landscape. we also found that carefully managing activation outliers is crucial to avoid performance drops with lower-precision quantization. additionally, we explain the sensitive nature of gradient quantization and its potential to fail at lower bit-widths. we observed that while the first-order moments of the adam optimizer can be effectively quantized to 4 bits, the second-order moments pose a greater challenge even for 8-bit quantization. overall, our work establishes an important foundation for future developments in quantization approaches tailored to transformers, opening the door to efficiently train large-scale models from scratch for improved accessibility.