Dynamic Stashing Quantization for Efficient Transformer Training (2023)

Abstract: Large Language Models (LLMs) have demonstrated impressive performance on a range of Natural Language Processing (NLP) tasks. Unfortunately, the immense amount of computations and memory accesses required for LLM training makes them prohibitively expensive in terms of hardware cost, and thus challenging to deploy in use cases such as on-device learning. In this paper, motivated by the observation that LLM training is memory-bound, we propose a novel dynamic quantization strategy, termed Dynamic Stashing Quantization (DSQ), that puts a special focus on reducing the memory operations, but also enjoys the other benefits of low precision training, such as the reduced arithmetic cost. We conduct a thorough study on two translation tasks (trained-from-scratch) and three classification tasks (fine-tuning). DSQ reduces the amount of arithmetic operations by 20.95\times and the number of DRAM operations by 2.55\times on IWSLT17 compared to the standard 16-bit fixed-point, which is widely used in on-device learning.

## introduction large language models (llms) based on the transformer architectures @xcite are currently seen as the foundation models @xcite . the pre-train and then finetune paradigm has shown promising results for a variety of natural language processing (nlp) tasks @xcite @xcite . however, the training of llm is both computationally and memory intensive, posing a significant challenge for their deployment. in the hardware world, the roofline model demonstrates that there is an optimal balance of processor and memory performance. the metric used to assess performance is referred to as the operational intensity, which is calculated as the ratio of arithmetic intensity to memory bandwidth: ## method figure 2 provides a high-level illustration of the dsq flow. we consider the inputs x l of a neural network layer with parameters w l , and the output of the layer is x l+1 . in the backward pass, we consider the partial derivatives dx l of the input and also the gradient of the weights dw l . naturally, a single training step requires three gemms as illustrated in figure 2 . we illustrate four quantization opportunities in this training step and their effects: in our knowledge, we are the first to systematically illustrate the potential effects, both on compute and off-chip memory bandwidth, of various quantization opportunities within a standard training pass. the two gemms in the backward pass can be potentially fused (e.g. pipelined), and in that case dx l does not have to be written to and then read from the dram. in our cost model estimation, we use a conservative strategy and assume this tensor is always flushed to dram. in dsq, we use block floating point as the quantizer for q 0 , q 1 , q 2 and q 3 , since this quantizer is shown superior to fixed-point quantization @xcite . we also use a time-adaptive quantization strategy, this means the quantization uses a different quantization level q t i for each round t of the training. we design dsq to monotonically increase q t i as a function of t and use the validation loss to inform this increase. this monotonic increase strategy has been proven more effective than other complex scheduling methods in @xcite . through extensive tuning and experimentation, we also notice that it is important to keep q 3 ≥ 16 through the entire training process, and appendix c studies the effect of different quantization levels for q 3 . ## evaluation we evaluate the effectiveness of dsq on two different translation tasks, wmt14 en-de @xcite (in appendix d) and iwslt17 en-de @xcite , and two tasks from the glue benchmark @xcite , the details of these datasets are in appendix a. we used the adam optimizer and the details for all the learning rate and batch size selections are in appendix b. for the translation tasks, we use a classic 6-layer transformer model @xcite and the roberta-base model @xcite for the glue tasks. all tasks are executed on systems that have 2 amd epyc 7763 64-core processors 1.8ghz (128 cores in total), and 4 nvidia a100-sxm-80gb gpus, with 1000 gib ram. we are interested in understanding the costs of arithmetic operations, as well as the number of memory reads and writes. to this end, we have built a hardware performance modeling framework to estimate the training cost. our cost model is similar to @xcite and @xcite , but our numbers are derived from a production hardware system, taking the numbers reported in darvish @xcite , to provide a higher-fidelity estimation. ## conclusion in this paper, we propose dynamic stashing quantization (dsq) for llm training. this new quantization strategy applies a more aggressive quantization for intermediate results between the forward and backward passes generated during training, thereby reducing dram traffic. specifically, our approach uses a low precision at the beginning of training, and then gradually increases the precision level, to reduce the effect of round-off errors introduced by quantization. we demonstrate the effectiveness of dsq by showing how it can reduce both the computation cost and dram bandwidth requirement on machine translation and llm fine-tuning tasks.