Dynamic Stashing Quantization for Efficient Transformer Training
Authors: Guo Yang
Daniel Lo
Robert D. Mullins
Yiren Zhao
Venue: EMNLP 2023 Findings (2023)
Keywords: Transformer
Training
Dynamic Stashing Quantization

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.

Dynamic Stashing Quantization for Efficient Transformer Training
Guo Yang
Robert D. Mullins
University of Cambridge Cambridge, UK gy2610cam. 10cam.cam..m ac uk
University of Cambridge
Cambridge, UK robert. www..t.mlllnnne..mmmmmmmmm.m cam ac .k
Daniel Lo Microsoft Research Redmond, Washington, USA
Yiren Zhao
Imperial College London London, UK
a. ww.zho@@meeiia..ccmm ac.
Abstract
Memory Bound
Compune Bound
Large Language Models (LLMs) have demon- strated impressive performance on range of Natural Language Processing (NLP) tasks. Un- fortunately, the immense amount of computa- tions and memory accesses required for LLM training makes them prohibitively expensive in terms of hardware cost, and thus challenging to deploy in cases such on-device learning.
Paffrrmne
Amaraabe Pefcrmmec
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 special focus on reducing the mem- ory operations, but also enjoys the other ben- efits of low precision training, such the re- duced arithmetic cost. We conduct a thorough study
Optina
Opaniival FEOR
Bashvill
hdantyy
Figure 1: The Roofline model with operational intensity (T) and attainable performance (P). 1 is non-quantized, 2 is a standard quantization and is DSQ.
translation tasks ttraiedd-frome scratch) and three classification tasks (fine- tuning). DSQ reduces the amount of arith- metic operations by 20.95x and the number of DRAM operations by 2.55 x on IWSLT17 compared to the standard 16-bit fixed-point.
Number of Operations
Operational Intensity
DRAM traffic
The Roofline model has enabled identify the sweet spot (Loyt) for processor to reach its peak arithmetic performance (Williams et al., 2009; Ding et al., 2022). As illustrated in Figure 1, as operational intensity (/) increases, the maximum attainable performance rises at linear rate initially before reaching constant value. The region to the left of the turning point is limited by the available memory bandwidth; the region to the right con- strained by the processor's arithmetic computing capability. Training Transformer models, as shown by Ivanov et al. (2021), is memory-bound. which means it sits at the left quadrant in the Roofline model (I < Lopt). Consequently, the performance of LLM training on modern hardware is signifi- cantly hindered by the inadequate bandwidth, the amount of data movements and from DRAM is the major performance bottleneck.
1 Introduction
Large Language Models (LLMs) based on the Transformer architectures (Vaswani et al., 2017) are currently
the foundation models (Bom- masani et al., 2021). The pre-train and then fine- tune paradigm has shown promising results for a variety of Natural Language Processing (NLP) tasks (Liu et al., 2019; Raffel et al., 2020; Brown et al., 2020). However, the training of LLM is both computationally and memory intensive, posing 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 referred to the op- erational intensity, which is calculated as the ratio arithmetic intensity to memory bandwidth:
For this reason, researchers have sought to accel- erate the training process of Transformers through

---

quantization. This approach aims reduce mem- ory consumption by decreasing the precision of parameters. Prior work has looked into the effect quantization on Transformer models, a majority which focus on the forward pass model infer- ence with fixed weights (Zhang et al., 2020; Bai et al., 2020; Tao et al., 2022). A number of studies have also investigated low low-precision training for Transformers (Sun et al., 2019, 2020). Although works have demonstrated the effectiveness of quan- tization, they typically assume single precision level, either per neural network layer or per net- work, which over-simplifies the hardware target. When viewed from Roofline model perspective. existing quantization methods attempt to optimize both compute complexity and memory bandwidth requirement, and then fail to recognize that the workload heavily memory-bound.
2.55x increase in arithmetic performance and a 20.95x reduction in DRAM requirement compared to 16-bit fixed-point training.
Related Work
Quantization has been studied in detail for infer- ence. These include using uniform (Zafrir et al., 2019; Bhandare et al., 2019) and non-uniform (Sun et al., 2019; Darvish Rouhani et al., 2020) quanti- zation methods. Specifically, uniform quantization methods such as fixed-point (Zafrir et al., 2019; Bhandare et al., 2019; Lin et al., 2020), ternary (Zhang et al., 2020), or even binary (Bai et al., 2020) number formats have been applied to infer- ence of Transformer models. In this work, we focus on quantization for LLM training which in- troduces new challenges such as the large dynamic range needed during the backward pass for loss- less training (Sun et al., 2019 99 where non-uniform quantization methods have seen more success.
Motivated by this observation, we propose a novel quantization strategy for LLM training named Dynamic Stashing Quantization (DSQ). We identify the most memory-intensin part of LLM training- - the communication between the forward and backward passes, and define stashing as the process of storing intermediate results in memory buffer (in a normal case, DRAM) for later use. The proposed quantization places an emphasis on this communication, and dynamically quantize the in- termediate results between forward and backward passes for a significant reduction of the DRAM traffic. As illustrated in Figure 1, this reduction of DRAM bandwidth helps DSQ to move closer to the optimal operational intensity. We have the following contributions:
Training LLM models is approximately 3x more expensive than running inference for the same model. Thus, quantizing all operations during train- ing has been an area of active research (Sun et al.. 2019, 2020; Yang et al., 2019; Fu et al., 2020; Fox et al., 2020; Kalamkar et al., 2019). Most of these methods use non- non-uniform quantization to handle larger dynamic range needed for gradient updates (Kalamkar et al., 2019). sooating-point arithmetic and its variants have become a popular method for low-precision training (e.g. fewer than 8 bits). Mini-floats with extremely small exponents (e.g. bit or bits) have been demonstrated to be effee- tive in small language models, such as LSTMs (Sun et al., 2019, 2020). Block floating-poin or block mini-floats, where exponent is shared between set of values, has become popular in quantized training (Yang et al., 2019; Drumond al., 2018; Fox et al., 2020) it allows for a large dynamic range while approximating the cost of integer for- mats for multiplication. Specifically, Draumond et al. utilized block floating-poin with roughly 24 bits to perform lossless training on vision tasks (Drumond et al., 2018). Fox et al. demonstrated that 8-bit training is possible with an around 0.5 BLEU score degradation machine translation (Fox al., 2020). Our work extends these formats to Large Language Models, includes quantization of stashed weights, and introduces a dynamic as- pect to further reduce the required bit widths. The idea of stashing has also been explored before by
We propose Dynamic Stashing Quantization (DSQ) for LLM training. DSQ not only quan- tizes operations for the entire training process, but also employs a more aggressive quanti- zation for intermediate results between the forward and backward passes to drastically minimize DRAM traffic.
DSQ follows a time-adaptive principle for stashing, which involves starting with lower precision at the beginning of the training pro- cess and gradually increasing the precision as it progresses. DSQ has been demonstrated to provide higher performance compared to its fixed -xecisccc counterpart.
We evaluate the proposed strategy on a variety of tasks and setups, including training from scratch and fine-tuning. DSQ achieves up to a

---

Jain et al. (2018), although they only focused on applying lossless encoding methods on single pre- cision numbers (6) However, in this paper we show a more aggressive stashing techniques (e.g. on average less than 4 bits per number) that is time-adaptive for LLM training. Fractrain (Fu et al., 2020), to our knowledge, is the only work that applied the idea of dynamic quantization on standard training, but was primarily focusing on vi- sion tasks. Our work extends dynamic quantization to encompass stashed values and evaluates these ef- fects on LLMs. Prior research on distributed train- ing has looked at reducing the communication cost (Alistarh et al., 2017; Hönig et al.. 2022), where Honig et al. also investigate how time-adaptive quantization would help federated systems to learn. These works focused on device-to to-ice traffic while our work focuses on reducing DRAM traffic,
this tensor is always flushed DRAM. In DSQ. we use Block Floating Point as the quantizer for qo- 91 92 and q3. since this quantizer is shown supe- rior to fixed-point quantization (Darvish Rouhani et al., 2020). We also use a time-adaptive quanti- zation strategy, this means the quantization uses different quantization level of for each round of the training. We design DSQ to monotonically increase of as function of t and use the valida- tion loss to inform this increase. This monotonic increase strategy has been proven more effective than other complex scheduling methods in Hönig et al. (2022). Through extensive tuning and ex- perimentation, we also notice that it is important to keep Q3 2 16 through the entire training pro- cess, and Appendix C studies the effect of different quantization levels for 43-
Evaluation
3 Method
We evaluate the effectiveness of DSQ on two dif- ferent translation tasks, WMT14 EN-DE (Bojar et al., 2014) (in Appendix D) and IWSLT17 EN- DE (Cettolo et al., 2017), and two tasks from the GLUE benchmark (Wang et al., 2018), 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 (Vaswani et al., 2017) and the RoBERTa-base model (Liu et al., 2019) 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 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 Sun et al. (2020) and Samajdar et al. (2018), but our numbers are derived from production hardware system, taking the numbers reported in Darvish Rouhani et al. (2020), to provide higher-fidelit estimation. Table presents the results of our study com- paring different quantization strategies. We com- pare popular low-latency training baselines and Block floating-point (BFP) (Darvish Rouhani al., 2020; Fox et al., 2020) with different precisions. For all BFP implementations considered in this paper, we keep the exponent bitwidth to be 8 and the bounding-box size to be 16 following
Figure 2 provides a high-level illustration of the DSQ flow. We consider the inputs xy of a neural network layer with parameters wi, and the output of the layer is T[+1- In the backward pass, we con- sider the partial derivatives dr of the input and also the gradient of the weights dui. Naturally, single training step requires three GEMMs as illus- trated in Figure 2. We illustrate four quantization opportunities in this training step and their effects:
Qo: mainly affects the arithmetic density of forward pass, notice it is possible for 2 and wi use different precisions, but this opti- mization is not the focus of our work. q1 affects the DRAM memory bandwidth, one key point in our work that we show @ can be different from Zo and in fact can be a very aggressive, dynamic quantization.
q2: affects mainly the computation complexity of the first GEMM in the backward pass. q3: affects the DRAM bandwidth and also the computation complexity of the second GEMM in the backward pass.
In our knowledge, we are the first to systemati- cally illustrate the potential effects, both on com- pute and off-chip memory bandwidth, of various quantization opportunities within a standard train- ing pass. The two GEMMs in the backward pass can be potentially fused (e.g. pipelined), and in that case da; does not have to be written to and then read from the DRAM. In our cost model estima- tion, we use a conservative strategy and assume

---

Frowird Paas
Badkward Pass
OEAM Clrma Maire Valler
GENM
--1--
GENM
dir
GEMM
dou
90
21
DRAM
DRAM
DSQ flow for single linear layer. The training is viewed combination forward pass and backward pass. go- g1 Q2 and @a define where the tensors quantized, we use (go. 91 22, Q3] to describe the DSQ configuration. DSQ ensures all GEMM inputs are quantized. Notice for the second and third GEMMs, dr 2i and dx, are the quantized version fetched from the DRAM, the fact that these values are heavily quantized helps us save DRAM bandwidth.
Figure 2:
An illustration
Table I:
The performance Machine Translation trained with 6-layer Transformer architecture, the model is assessed using numbers reported as percentages. A shows the performance difference compared to the floating-point 32-bit baseline.
Dataset and Model
Method
Precision Setup
Acc BLEU (A)
Arith Ops (1)
DRAM R/W (L)
Flosting-point
[32, 32, 32, 32]
36.22
Fixed point
[32, 32, 32, 32]
34.47 (-0.75)
1.00x
1.00x
Fixed-point
[16, 16, 16, 16]
32.59 (-2.63)
0.25x
0,50x
IWSLT2017 DE-EN Transformer (6-layer)
FP
[3,,32, 32,
34.56 0.66)
0.56x
Block FP
16,
31.00 (-0.2)
0.18x
0.63
Stashing (Fixed)
[16, 4,4, 16]
25.50 9.72)
0.13x
0.31
Stashing (BFP)
[16, 4.44 16
34.88 (-0.44)
0.10x
0.45
DSQ (BFP)
34.81 (-0.41)
0.012x
0.20x
Plsstnng-ooit
[32, 32, 32, 32]
87.6
Fixed Fixed-point
[3,, 32, 32]
87.9 (+0.3)
1.00x
1.00x
Fixed-point
[16, 16, 16]
87.9 (+0.3)
0.25x
0.50
GLUE MNLI
Block TP
[32, 32, 321
87.8 (10.2)
0.56x
1.13x
RoBERTa-base
Block FP
[16, 16, 16. 16]
87.8 (+0.2)
0.18x
0.63
Stashing (Fixed)
[16. 4,
82.8
0.32x
Stashing (BFP)
4,4,4,4,16
(+0.2)
0.10x
0.45x
DSQ (BFP)
87.8 (+0.2)
0.043x
0.26x
Floating point
[32, 32,
Fixed-poins poiit
[32, 32,
92.6 (-0.)
1.00
1.00
Fixed-point
[16, 16. 16]
92.6 (-0.)
0.25x
0.50x
GLUE QNLI RoBEKIs-base
Block FP
[32. 32. 321
92.7 (-0.)
0.56x
1. 1.13x
Block HP
[16, 16, 16]
92.5 (-033
0.18x
0.63%
Stushing (Fixed)
4,4,
89.5 (-33)
0.13x
0.32x
Stashing (BFP)
4,4,4,41
92.6 (-02)
0.10x
0.45x
DSQ (BFP)
92.7 0.1)
0.043x
0.26x
Darvish Rouhani et al. (2020). In addition, compare static stashing strategies that are based either fixed-point (Fixed) BFP. In Table 1, use the hardware cost of fixed-point 32-bit computation as 1x since this is a stronger bascline. The results demonstrate that DSQ has a comparable accuracy and BLEU score compared to 16-bit fixed-point for training while having 20.95x reduction in arith- metic complexity and a 2.55x decrease in DRAM R/W. DSQ also shows very competitive accuracy on fine-tuning RoBERta on GLUE while having much smaller hardware utilization.
5 Conclusion
In this paper, we propose Dynamic Stashing Quanti- zation (DSQ) for LLM training. This new quantiza- tion 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 com- putation cost and DRAM bandwidth requirement on machine translation and LLM fine-tuning tasks.

---

6 Limitation
Rishi Bommasani, Drew Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg. Antoine Bosse- lut, Emma Brunskill, al. 2021. On the opportuni- ties and risks of foundation models. arXiv preprint
DSQ precision configurations are decided through experimentation on the IWSLT dataset. The precision for different stages is scheduled based on the validation loss value, the precision would increase the validation loss becomes 'flat' (non- (non-decreasing). In our particular method, if our validation loss has been non-decreasing for a fixed number of N epochs, we then move to the next quantiza- tion level, following a setup similar to that proposed by Hönig et al. (2022). We find that setting N = 5 is sufficient for all test sce- narios. The choice of this free parameter is a limitation in our paper and will be investigated in further research. The same precision con- figuration setup is used for all other datasets.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901.
Mauro Cettolo, Marcello Federico, Luisa Bentivogli. Niehues Jan, Stüker Sebastian, Sudoh Katsuitho. Yoshino Koichiro, and Federmann Christian. 2017. Overview of the iwslt 2017 evaluation campaign. In International Workshop on Spoken Language Trans- lation, pages 2-14.
Bita Darvish Rouhani, Daniel Lo, Ritchie Zhao, Ming Liu, Jeremy Fowers, Kalin Ovtcharov, Anna Vino- gradsky. Sarah Massengill, Lita Yang. Ray Bittner. et al. 2020. Pushing the limits of narrow precision in- ferencing at cloud scale with microsoft floating point. Advances in neural information processing systems, 33: 10271-00221
Language models that are larger than classic Transformer and RoBERTa have been devel- oped in recent works. Due to limited resouces we have, we choose to work on smaller mod- els as an exemplar. While we have tested on larger LMs like OPT-1.3B (in Appendix E) to show large L.Ms today are also memory- bound, additional experimentation is expected to be conducted to enhance the robustness and precision of our findings.
Nan Ding, Muaaz Awan, and Samuel Williams. 2022. Instruction roofline: An insightful visual perfor- mance model for gpus. Concurrency and Compu- lation: Practice and Experience, 34(20):66591.
Mario Drumond, Tao Lin, Martin Jaggi, and Babak Fal- safi. 2018. Training dans with hybrid block floating point. Advances in Neural Information Processing Systems, 31.
References
Sean Fox, Seyedramin Rasoulinezhad, Julian Faraone, Philip Leong. et al. 2020. A block minifloat rep- resentation for training deep neural networks. In International Conference on Learning Representa- tions.
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovie. 2017. Qsgd: Communication- efficient sgd via gradient quantization and encoding. Advances in neural information processing systems, 30.
Haoli Bai, Wei Zhang, Lu Hou, Lifeng Shang, Jing Jin, Xin Jiang, Qun Liu, Michael Lyu, and Irwin King. 2020. Binarybert: Pushing the limit of bert quantization. arXiv preprint arXiv:2012. 15701.
Yonggan Fu, Haoran You, Yang Zhao, Yue Wang, Chao- jian Li, Kailash Gopalakrishnan, Zhangyang Wang, and Yingyan Lin. 2020. Fractrain: Fractionally squeezing bit savings both temporally and spatially for efficient dnn training. Advances in Neural Infor- mation Processing Systems, 33:12127-12139,
Aishwarya Bhandare, Vamsi Sripathi, Deepthi Karkada, Vivek Menon, Sun Choi, Kushal Datta, and Vikram Saletore. 2019. Efficient& 8-bit quantization of trans- former neural machine language translation model. arXiv preprint arXiv: 1906.00532.
Robert Hönig, Yiren Zhao, and Robert Mullins. 2022. Dadaquant: Doubly-adaptive quantization for cummunication-effiiie federated learning. In In- ternational Conference on Machine Learning, pages 8852-8866. PMLR
Ondrcj Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz. Pavel Pecina, Matt Post, Herve Saint- Amand, et al. 2014. Findings of the 2014 4 workshop on statistical machine translation. In Proceedings of the ninth workshop on statistical machine translation, pages 12-58.
Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. 2021. Data movement is all you need: A case study optimizing transform- ers. Proceedings of Machine Learning and Systems. 3:711-732.

---

Animesh Jain, Amar Phanishayee, Jason Mars, Lingiaa Tang, and Gennady Pekhimenko. 2018. Gist: Effi cient data encoding for deep neural network train- ing. In 2018 ACM/IEEE 4Sth Annual International Symposiam on Computer Architecture (ISCA), pages 776-789. IEEE
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. 2018. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 1vXv::880.007411
Samuel Williams, Andrew Waterman, and David Pat- terson. 2009, Roofline: an insightful visual perfor- mance model for multicore architectures. Communi- cations of the ACM, 52(4):65-76.
Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellem- pudi, Dipankar Das, Kunal Banerjee, Sasikanth Avan-
cha, Dharma Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang. Hector Yuen, et al. 2019. A study of 6 for deep learning training. arXiv preprint arXiv:1905. 12322..
Guandao Yang, Tianyi Zhang, Polina Kirichenko, Jun- wen Bai, Andrew Gordon Wilson, and Chris De Sa. 2019. Swalp: Stochastic weight averaging in low precision training. In International Conference on Machine Learning, pages 7015-7024. PMLR.
Ye Lin, Yanyang Li, Tengbo Liu, Tong Xiao, Tongran Liu, and Jingbo Zhu. 2020. Towards fully 8-bit in- teger inference for the transformer model. arXiv preprint KXi::20090080344
Ofrr Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. 2019. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing -NerrIPS Edition (EMC2-NIPS), pages 36-39, IEEE
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Man- dar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019.
Roberta: A robustly optimized bert pretraining ap- proach. arXiv preprint arXiv: arXiv:1907.1 1et.. 1692.
Wei Zhang. Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, and Qun Liu. 2020. Ternarybert: Distillation-awall ultra-low bit bert. arXiv preprint arXiv:2009.12812.
Colin Raffel, Noam Shazeer. Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter Liu, et al. 2020. Exploring the limits of transfer learning with unified text-to-text trans- former. J. Mach. Learn. Res., 21(140):1-67.
A Datasets
Four datasets are used: translation WMT14 EN-DE and IWSLT2017 EN-DE for machine translation tasks, QNLI and MNLI for textual entailment tasks. Table 2 presents details for the datasets.
Ananda Samajdar, Yuhao Zhu, Paul Whatmough, Matthew Mattina, and Tushar Krishna. 2018. Scale- sim: Systolic cnn accelcraoor simulator. arXiv preprint arXiv: 1.28883.
B Hyperparameters
Xiao Sun, Jungwook Choi, Chia-Yu Chen, Naigang Wang, Swagath Venkataramani, Vijayalakshmi Viji Srinivasan, Xiaodong Cui, Wei Zhang, and Kailash Gopalakrishnan. 2019. Hybrid 8-bit floating point (hfp8) training and inference for deep neural net- works. Advances in neural information processing systems, 32.
The training hyperparamters, such as learning rates, are picked following standard benchmarks and open implementaitons (Liu al., 2019; Vaswani et al., 2017). We summarize them in Table 3 for repeatability. We use the Adam optimizer with B1 -1-0.9 0.9, B2 0.98 for both training and finetun- ing models. The learning rate schedule is Inverse Square Root for training models. and Polynomial Decay for finetuning models. Dropout with rates of PIWSLT 0.3 and Pw MT 0.2, label smooth- ing with value e 0.1 are applied to train models.
Xiao Sun, Naigang Wang. Chia- Yu Chen, Jiamin Ni, Ankur Agrawal, Xiaodong Cui, Swagath Venkatara- mani, Kaoutar El Maghraoui, Vijayalakshmi Viji Srinivasan, and Kailash Gopalakrishnan. 2020. Ultra- low precision 4-bit training deep neural networks. Advances in Newral Information Processing Systems, 33:1796-1807
DSQ precision configurations are decided through experimentation on the IWSLT dataset as discussed in section 6 Limitations. Table 4 shows collection of tuning runs we had, we found that heavily quantized models still work at the start of training stage, and [16, 4, 4, 16] quantized BFP model works as well as less aggressive ones. This indicates that DSQ should start with heavily ag- gressive precision setup (we pick [2, 2, 2, 16] for IWSLT14 DE-EN), and jump to [16, 4, 4, 16] when needed during training process.
Chaofan Tao, Lu Hou, Wei Zhang, Lifeng Shang, Xin Jiang, Qun Liu, Ping Luo, and Ngai Wong. 2022. Compression of generative pre-trained lan- guage models via quantization. arXiv preprint arXiv:2203. 10705.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural procesnnn systems, 30.
processing

---

Table 2: Details for each dataset, including the number of classes, description and the souree.
Name
# Class
Description
A text translation task on English-Germane sentence pairs from
WMT14 EN-DE
The The Stanford Natural Language Processing Group.
A text translation task on German German-English sentence pairs from The International Conference Spoken Language Translation.
IWSLT2017 DE-EN
A binary lextual entailment lask on question-ana -nnswe paits from the Stanford Question Answering database. The objective is to determine whether pair is an entailment
QNLI
A multi-class (i.e.. entailment, neutral, contradiction) textual en- tailment task premise- premise-hypothesie pairs from the Multi-genre Natural Language Inference corpus, Matched version only pre- serves pairs within the same genre (e.g., government report, sci- ence fiction, speech).
MNLI
Table 3: Details of the optimal hyper- param-aaaseterss including batch size, learning rate and weight decay values for each set of experiments with the same dataset and prompting model.
Dataset
Batch size
Max tokens
Learning rate
Weight decay
WMT14 EN-DE
4096
5e-4
0.0
IWSL12017 DE-EN
4096
5e4
lc-4
QNLI MNLI
4400
lc-5
0.1
4400
Table Tests on stashing precision setup. The models are trained IWSLT14 DE-EN. A shows the performance difference compared to the doating-point 32-bit baseline.
Datuset and Model
Method
Precision Setup
Acc BLEU (A)
Stushing (BFP)
[2, 16]
17.45 (-17.77)
Stashing (BPP)
[4, 16]
33.51 (-1.71)
Stashing (BI+P)
[4,4, 16]
34.47 (-U.75)
IWSLT14 DE-EN
Stashing (BIFP)
18,4.4,
31.17 (-0.75)
Transformer (6-layer)
Stashing (BFP)
[8, 16]
34.65 (-0.57)
Stashing (BFP)
[16, 4, 4, 16]
34.78 (-0.44)
Stashing (BFP)
[16, 8, 8, 16]
34.47 (-0.75)
The effect of Q3
E Memory-bound verification on OPT-1.3B
The gradient output (dxz) plays an important role in the performance of fixed-poin quantization. Notice in table Table 5, gradient output quantized to 8 bits leads to training failure for fixed-pond point quantization. In order to focus the idea of stashing, apply 16 bits quantization of gradient output for all our stashing precision setups.
we have also tested on larger LMs like OPT-1.3B to show large LLMs today are also memory-bound so that dynamic stashing will serve as a useful tech- nique for reducing the DRAM traffic and compute. We run a similar amalysis to Ivanov et al. (2021). Table 7 presents the results. The analysis agrees with the conclusion of Ivanov et al. (2021). This may be of interest the community in understand- ing the memory-bound nature of current LLMs.
D Additional results on WMT14
We also train the model WMT14 EN-DE dataset, the BLEU scores we gain are relatively low com- pared to the 27.3 BLEU score achieved by Vaswani et al. (2017) because we only trained the models for 15 epochs. Table 6 presents the results.

---

Table 5: Tests on gradient output precision setup. The models are trained on IWSLT14 DE-EN.
Dataset and Model
Method
Precision Setup
Acc/ BLEU (A)
Stashing (Fixed)
[8, 8, 8, 32]
34.08
4 DE-EN
Stushing (Fixeed
[8. 8, 8, 16]
31.94
Transformer (6-layer)
Stushing (Fixed)
[8, 8,
Failed
Table 6: The performance of Machine Translation trained WMT14 EN-DE with 6-layer Transformer architecture (Vaswani et al., 2017), the model is assessed using numbers reported as percentages. A shows the performance difference compared to the floating-pote 32-bit baseline.
Dataset and Model
Method
Precision Sctup
Acc BLEU (A)
Arith Ops
DRAM R/W
Floating-point
[32, 32, 32, 32]
25.79
Fixed-point
[32, 32, 32, 32]
25.41 (-0.38)
1.00x
1.00x
Fixed-point
[16, 16, 16, 16]
23.40 (-2.39)
0.25x
0.50x
WMT14 EN-DE
Block FP
[32, 32,
25.76 (-0.03)
0.56x
1.13x
Transformer (6-layer)
Block FP
[16, 16, 16, 16]
25.61 (-0.18)
0.18x
0.63x
Stushing (Fixed)
[16, 4, 4. 16]
21,86 3.93)
0.13x
0.31
Stashing (BFP)
[16, 4, 16]
25.24 00.55)
0.10x
0.20x
Table 7: Bound type analysis of OPT-1.3B based computational complexity and latency fraction
Model
Attention Layers
Layers
Others
Latency Praction
47.93%
32.20%
19.87%
OPT-113R
Bound Type
Memory-bound
Compute-bound
Memory-bound


Reviews:
[{"review_id": "6j9eFCrZHw", "reviewer": "Reviewer_jF2i", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments.\n\nIt is worth noting that there is a growing trend towards utilizing domain-specific hardware for accelerating ML workloads, exemplified by Google’s TPU, AMD’s AI Engine and Cerebras WSE for AI acceleration as compared to NVIDIA’s GPGPU. In this regard, we hold the belief that ASICs optimized for efficiency and cost will be the future generation hardware for language model training.\n\nNVIDIA’s GPU is the most common choice for training and testing neural networks in the deep learning community because deep learning frameworks like PyTorch have high-performance tensor computation libraries built on CUDA e.g., Aten for PyTorch. However, one limitation of GPUs is the lack of native support for custom quantization arithmetics like the BFP mentioned in this paper. For example, NVIDIA RTX 4090 only supports FP64, FP32, Bfloat16, FP16, and Int8, and PyTorch supports training & testing with FP64, FP32, Bfloat16, and FP16. Though sometimes it is possible to support custom quantization arithmetics via manually-designed CUDA kernel, the efficiency is still much lower than ASICs.\n\nIn conclusion, we are exploring next-generation hardware for training language models, and the implementation of DSQ requires specialised hardware. Potential challenges may include hardware support and existing algorithm adaptation. We will add this explanation to our revised version of the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The concept of Dynamic Stashing Quantization (DSQ) is innovative and its emphasis on aggressive quantization for intermediate results is well-motivated. The time-adaptive stashing principle introduces a dynamic element to the quantization process, which is a notable contribution.\n\nThe experimental evaluation is thorough, encompassing various tasks and scenarios, and showcasing DSQ's superiority over fixed-precision methods. The achieved performance improvements are impressive, with up to a 2.55× boost in arithmetic performance and a substantial 20.95× reduction in DRAM requirements compared to 16-bit fixed-point training.\n\n\n\n", "reasons_to_reject": "While the paper's content is solid, I believe some sections could benefit from further elaboration. Specifically, more details on the practical implementation of DSQ, as well as potential challenges and limitations, would enhance the paper's completeness.", "questions_to_authors": "", "paper_topic_and_main_contributions": "Dynamic Stashing Quantization (DSQ) is introduced as a novel approach for Language Model (LLM) training. It involves quantizing operations throughout training and emphasizes aggressive quantization for intermediate results between forward and backward passes to minimize DRAM traffic. DSQ employs a time-adaptive stashing principle, starting with lower precision and gradually increasing it during training. This technique outperforms fixed-precision methods, resulting in significant improvements in performance. Evaluation across various tasks and setups, including training from scratch and fine-tuning, demonstrates DSQ's effectiveness. It achieves up to a 2.55× boost in arithmetic performance and reduces DRAM requirements by 20.95× compared to 16-bit fixed-point training.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "None"}, "extra_text": {}}, {"review_id": "t624hzJvYZ", "reviewer": "Reviewer_vpzW", "summary": "", "questions": "1. What is the time to train for the baseline models evaluated in this paper?\n2. How does this technique interact with gradient checkpointing, replays and gradient accumulation? Standard techniques used from LLM training. ", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments.\n\nThe compute resource is a high-performance computing facility that is shared among a number of researchers. Due to the limited compute hours available to us, we have chosen to work on smaller models as an exemplar. Nevertheless, in Appendix E, We run a similar analysis to Ivanov et al. on larger LMs like OPT-1.3B to show LLMs today are also memory-bound, so that DSQ will serve as a useful technique for reducing the DRAM traffic and compute.\n\nOur answers to the reviewer's questions:\n\nQuestion 1:\n\nIt is worth noting that there is a growing trend towards utilizing domain-specific hardware for accelerating ML workloads, exemplified by Google’s TPU for AI acceleration as compared to NVIDIA’s GPGPU. In this regard, we hold the belief that ASICs optimized for efficiency and cost will be the future generation hardware for language model training.\n\nNVIDIA’s GPU is the most common choice for training and testing neural networks in the deep learning community because deep learning frameworks like PyTorch have high-performance tensor computation libraries built on CUDA e.g., Aten for PyTorch. However, one limitation of GPUs is the lack of native support for custom quantization arithmetics like the BFP mentioned in this paper. For example, NVIDIA RTX 4090 only supports FP64, FP32, Bfloat16, FP16, and Int8, and PyTorch supports training & testing with FP64, FP32, Bfloat16, and FP16. Though sometimes it is possible to support custom quantization arithmetics via manually-designed CUDA kernel, the efficiency is still much lower than ASICs.\n\nIn conclusion, we are exploring next-generation hardware for training language models, which requires the implementation of DSQ on specialised hardware in order to achieve a speed-up. We will revise our draft to include this to our Appendix.\n\nQuestion 2:\n\nDSQ operates as a quantization methodology and does not inherently change the way of using gradient checkpointing, replays and gradient accumulation. In other words, these techniques can be applied in conjunction to DSQ. However, it is worth noting that the precision of DSQ changes during training, and makes less utilisation of the memory. In this case, if applying gradient checkpointing/replay, we should in theory see less layers being recomputed in the backward pass. Moreover, this method operates in a separate optimization space when it comes to gradient accumulation.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Paper is easy to ready and understand\n2. The focus on quantization of gradient communication between forward and backward is interesting. ", "reasons_to_reject": "1. The method while interesting needs a lot more benchmarking on more complex tasks to ensure that it works generally. Pushing for a larger model like BERT will also help. For finetuning tasks like SQuAD will also add credence to the work. \n\n2. Without any runtime numbers, its hard to understand the wall clock impact of these techniques. Additionally, arithmetic Ops reduction will not necessary translate to faster runtime for arbitrary quantization bit-widths. As a result, the impact of this work may be limited.", "questions_to_authors": "1. What is the time to train for the baseline models evaluated in this paper?\n2. How does this technique interact with gradient checkpointing, replays and gradient accumulation? Standard techniques used from LLM training. ", "paper_topic_and_main_contributions": "This paper focuses on a dynamic quantization algorithm thats majorly focussed on reducing the DRAM traffic during training and hence improve the operational intensity of the problem. They improve upon the usual quantization flow by identifying that the quantization of weights and activations during the forward and backward can be different. Especially, the activations being read for gradient calculations can be significantly quantized. The algorithm focusses on 4 different quantization operations - activations and weights during forward calculation, output activation after forward call being written to DRAM for gradient calculation, activations being read for gradient calculations and weight quantization during gradient calculation. The quantization level changes as training progresses. ", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "1. There has been a lot of work in the domain of quantization of gradients used for synchronization across DP workers. The work in this domain should go in related work section"}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "1. A nit, talking about LLMs in Introduction creates a wrong impression of the problem being solved in the paper what to expect during the evaluation section. 6 layer roberta should not be constituted as LLMs. "}}, {"review_id": "I2vesSsjUV", "reviewer": "Reviewer_3Ppe", "summary": "", "questions": "How does DSQ compare to other existing quantization strategies in terms of efficiency and performance?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments.\n\nThe compute resource is a high-performance computing facility that is shared among a number of researchers. Due to the limited compute hours available to us, we have chosen to work on smaller models as an exemplar. Nevertheless, in Appendix E, We run a similar analysis to Ivanov et al. on larger LMs like OPT-1.3B to show LLMs today are also memory-bound, so that DSQ will serve as a useful technique for reducing the DRAM traffic and compute.\n\nIt is worth noting that there is a growing trend towards utilizing domain-specific hardware for accelerating ML workloads, exemplified by Google’s TPU for AI acceleration as compared to NVIDIA’s GPGPU. In this regard, we hold the belief that ASICs optimized for efficiency and cost will be the future generation hardware for language model training.\n\nNVIDIA’s GPU is the most common choice for training and testing neural networks in the deep learning community because deep learning frameworks like PyTorch have high-performance tensor computation libraries built on CUDA e.g., Aten for PyTorch. However, one limitation of GPUs is the lack of native support for custom quantization arithmetics like the BFP mentioned in this paper. For example, NVIDIA RTX 4090 only supports FP64, FP32, Bfloat16, FP16, and Int8, and PyTorch supports training & testing with FP64, FP32, Bfloat16, and FP16. Though sometimes it is possible to support custom quantization arithmetics via manually-designed CUDA kernel, the efficiency is still much lower than ASICs.\n\nIn conclusion, we are exploring next-generation hardware for training language models, which requires the implementation of DSQ on specialised hardware in order to achieve a speed-up. We will revise our draft to include this to our Appendix.\n\nOur answer to the reviewer's question:\n\nWe noticed that DRAM bandwidth is a more critical limiting factor for Transformer training than for training vision networks. Most existing quantization methods treat all components in the network equally and utilise mainly fixed-point quantization. To combat this, we implemented both fixed-point quantization and BFP quantization, and found out the BFP performed better in Transformer setup. We then proposed a dynamic (time-adaptive) stashing technique with BFP (DSQ) to reduce DRAM traffic. The precision of DSQ can go up to 32 bits, providing us with a nice fall-back scenario where the accuracy can match that of large bitwidth quantization methods.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper tackles an important problem in the realm of LLM training, namely the high computational and memory requirements.\n2. The idea of time-adaptive compression, which involves starting with lower precision at the beginning of the training process and gradually increasing the precision as it progresses, is interesting and novel.\n3. The authors provide a sufficient evaluation of DSQ, including tests on both translation and classification tasks. The evaluation adds to the paper's credibility and provides a robust demonstration of DSQ's effectiveness.", "reasons_to_reject": "1. The paper primarily focuses on smaller-scale models, leaving some uncertainty about the applicability and effectiveness of the proposed DSQ method on more complex, large-scale models. It would significantly enhance the paper's value if the authors could extend their experiments to include fine-tuning of larger models, such as OPT 1.3B and LLaMA, and provide a comparative analysis of the perplexity outcomes.\n2. The paper provides a comparison of arithmetic complexity and DRAM Read and Writes, but it would be great to include an end-to-end speedup analysis, which would offer a more holistic understanding of DSQ's performance in real-world scenarios.", "questions_to_authors": "How does DSQ compare to other existing quantization strategies in terms of efficiency and performance?", "paper_topic_and_main_contributions": "This paper addresses the challenges associated with training Large Language Models (LLMs), particularly focusing on the high computational and memory demands that make such models expensive and difficult to deploy in situations like on-device learning. The authors propose Dynamic Stashing Quantization (DSQ), a novel dynamic quantization strategy designed to significantly reduce memory operations while also benefiting from the reduced arithmetic cost associated with low precision training. The method's efficacy is demonstrated through a comprehensive study on two translation tasks (trained-from-scratch) and three classification tasks (fine-tuning).", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hCzbzRphHq", "reviewer": "Reviewer_Qa9X", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments.\n\nThe compute resource is a high-performance computing facility that is shared among a number of researchers. Due to the limited compute hours available to us, we have chosen to work on smaller models as an exemplar. Nevertheless, in Appendix E, We run a similar analysis to Ivanov et al. on larger LMs like OPT-1.3B to show LLMs today are also memory-bound, so that DSQ will serve as a useful technique for reducing the DRAM traffic and compute.\n\nWe thank the reviewer for the comments regarding training stability. Empirically, we observe that our training log figures show stable gradients and loss. We will include the figures in our final draft.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1) The motivation is great. Reducing memory consumption and FLOPs is very important for LLM training efficiency.\n2) The results are very impressive. It is good to write such a concise idea in a short paper. ", "reasons_to_reject": "1) The model is too small. I understand the resource is a problem. Authors do not need to fix this during rebuttal. But I have to say this as a limitation.\n2) Are there any training stability issue? Since the model used is small, the training stability is not a big issue. But I believe it would be beneficial to try a larger model with higher learning rate to check how stable the approach is.", "questions_to_authors": "", "paper_topic_and_main_contributions": "This paper proposes a novel dynamic quantization strategy, termed Dynamic Stashing Quantization (DSQ) to reduce the amount of computations and memory accesses required for LLM training.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: When I look at this paper, I see some flaws and statements that will not translate to real hardware and personally, I find it a bit misleading. However, the overall response from the reviewers is strongly in factor of acceptance both in soundness and excitement, and I will not overrule their decision. As such, I recommend acceptance to the main conference.
Decision: Accept-Findings