Curriculum Learning: A Regularization Method for Efficient and Stable Billion-Scale GPT Model Pre-Training

Recent works have demonstrated great success in training high-capacity autoregressive language models ( GPT , GPT-2 , GPT-3 ) on a huge amount of unlabeled text corpus for text generation . Despite showing great results , autoregressive models are facing a growing training instability issue . Our study on GPT-2 models ( 117M and 1.5B parameters ) show that larger model sizes , sequence lengths , batch sizes , and learning rates would lead to lower training stability and increasing divergence risks . To avoid divergence and achieve better generalization performance , one has to train with smaller batch sizes and learning rates , which leads to worse training efficiency and longer training time . To overcome this stability-efficiency dilemma , we present a study of a curriculum learning-based approach , which improves the pre-training convergence speed of autoregressive models . More importantly , we find that curriculum learning , as a regularization method , exerts a gradient variance reduction effect and enables to train autoregressive models with much larger batch sizes and learning rates without training instability , further improving the training speed . Our evaluations demonstrate that curriculum learning enables training GPT-2 models with 8x larger batch size and 4x larger learning rate , whereas the baseline approach struggles with training divergence . To achieve the same validation perplexity targets during pre-training , curriculum learning reduces the required number of tokens and wall clock time by up to 61 % and 49 % , respectively . To achieve the same or better zero-shot WikiText-103/LAMBADA evaluation results at the end of pre-training , curriculum learning reduces the required number of tokens and wall clock time by up to 54 % and 70 % , respectively . 1 INTRODUCTION . Large-scale Transformer-based language models have powered breakthroughs in many natural language processing tasks ( Vaswani et al. , 2017 ) . These models contain up to billion-scale parameters and are trained on massive open-domain web text corpus with unsupervised learning . Among them , one of the most successful cases is the GPT family : from GPT ( 110M parameters ) ( Radford et al. , 2018a ) , to GPT-2 ( 1.5B parameters ) ( Radford et al. , 2018b ) , and to GPT-3 with 175B parameters ( Brown et al. , 2020 ) that hit record high accuracy for many NLP tasks . Recent studies ( Kaplan et al. , 2020 ) show that these models continue to achieve better accuracy as their sizes increase ( together with richer training data ) , indicating the model size will continue growing in the future . Despite achieving remarkable model accuracy , training these models raise huge challenges on training efficiency and instability . On one hand , training large models require a huge amount of computation . As an example , it requires approximately 9.2 days on 512 V100 GPUs to train a 8.3B GPT-2 ( Shoeybi et al. , 2019 ) , and 14.8 days on 10000 V100 GPUs to train a 175B GPT-3 ( Patterson et al. , 2021 ) . To reduce the training wall clock time of these large models , one of the most common solutions is to employ distributed training with hundreds or even thousands of GPUs across multiple machines such that the model can process the training data with much higher throughput . Meanwhile , to avoid the huge overhead incurred by cross-GPU/node communication in distributed training , practitioners increase the batch sizes in the optimization algorithm to increase the computation-communication ratio , to improve training efficiency and reduce training wall clock time . Despite increased training throughput and efficiency , increasing the batch size is not always the panacea , as it often leads to training instability during pre-training GPT models — training would diverge or recover from divergence but overall lead to slow/worse convergence . We conduct a thorough study of the GPT-2 pre-training task ( Radford et al. , 2018b ; Shoeybi et al. , 2019 ) with different batch sizes , models sizes , sequence lengths , and learning rates . We find that larger batch sizes and longer sequence lengths reduce training stability and increase divergence risk , especially when combined with larger model sizes and learning rate . This often leads to slower and worse convergence that hurts final generalization performance . To achieve stable training and avoid suboptimal generalization performance , common practice often resorts to using smaller batch sizes and learning rates , which would adversely affect training efficiency . In this work , we explore a curriculum learning-based approach to address the above training efficiencystability dilemma . Curriculum learning ( CL ) was proposed to improve training convergence speed by presenting easier/simpler examples earlier during training and gradually increasing the sample difficulties ( Elman , 1993 ; Sanger , 1994 ; Bengio et al. , 2009 ) . CL was explored and verified for NLP one-stage and fine-tuning tasks ( Kocmi & Bojar , 2017 ; Bojar et al. , 2017 ; Zhang et al. , 2018 ; Platanios et al. , 2019 ; Zhang et al. , 2019 ; Sachan & Xing , 2016 ; 2018 ; Tay et al. , 2019 ; Xu et al. , 2020 ) , but its application to pre-training GPT models is not well studied . We find that curriculum learning , as expected , could improve the convergence speed under the same training hyperparameters . More importantly , we find that curriculum learning , as a regularization method , can enable stable and efficient training with much larger batch sizes and learning rates than existing baseline approaches . Our analysis shows that one reason for curriculum learning ’ s better training stability lies in its gradient variance reduction effect : Although the gradient variance is generally reduced under larger batch size , the largest variance on certain dimensions ( i.e. , the outliers ) is increased and it reaches extreme values when training divergence happens . Using Adam optimizer ’ s variance term , we observe that curriculum learning helps reduce both the norm of this variance and the maximum variance outliers . This variance reduction effect helps curriculum learning to achieve stable large-batch training without hurting the token-wise convergence speed and final generalization performance . We implement curriculum learning using sequence length as a primary difficulty metric and apply it to the GPT-2 pre-training with up to 1.5B parameters . Evaluations show that curriculum learning enables stable and efficient training with 8x larger batch size and 4x larger learning rate , where the baseline approach struggles with training divergence . To achieve the same validation perplexity targets during pre-training , curriculum learning reduces the required number of tokens and wall clock time by up to 61 % and 49 % , respectively . To achieve the same or better zero-shot WikiText-103/LAMBADA evaluation results at the end of pre-training , curriculum learning reduces the required number of tokens and wall clock time by up to 54 % and 70 % , respectively . We make the following contributions : i ) We conduct an extensive study of the GPT-2 pre-training task , which provides detailed insights about the training stability-efficiency dilemma that motivate our work ( Section 3 ) . ii ) We present an implementation of curriculum learning based on sequence length for GPT-2 model ( and autoregressive model in general ) , which is both efficient and easy to integrate ( Section 4 ) . iii ) We conduct large-scale experiments to demonstrate the proposed work ’ s ability to provide superior training stability and efficiency at the same time . To the best of our knowledge , we are the first work to demonstrate the benefit of curriculum learning as a regularization method that improves training stability ( Section 5 ) . iv ) The curriculum learning implementation as well as the necessary changes to the GPT-2 pre-training framework has been open sourced in a deep learning optimization library ( name hidden to maintain anonymity ) . 2 BACKGROUND AND RELATED WORK . Language model pre-training The accuracy of transformer-based language models grows substantially with its model size ( Radford et al. , 2018a ; b ; Brown et al. , 2020 ) . Today , a large language model such as GPT-3 ( Brown et al. , 2020 ) contains up to 175B parameters , and recent studies show that model accuracy can continue to improve with even larger model sizes ( Kaplan et al. , 2020 ) . However , training these large models often incurs excessively long training time and training difficulties ( Brown et al. , 2020 ) . Therefore , there are a lot of demands of performing efficient and stable training for large-scale LMs . To have the pre-training finished in a reasonable amount of time , the most common way is to leverage data parallelism to train models on multiple GPUs . However , the speedup gains often saturate beyond a few tens of GPUs , because communication becomes the major bottleneck , i.e. , the workers will spend more time communicating gradients than computing them , as the number of GPUs increases . To mitigate this bottleneck , recent works such as 1-bit Adam ( Tang et al. , 2021 ) have studied gradient compression and demonstrate their effectiveness against auto-encoding models such as BERT ( Devlin et al. , 2019 ) . An alternative approach to alleviate these overheads is to use large batch sizes . For example , LAMB ( You et al. , 2020 ) and 1-bit LAMB ( Li et al. , 2021 ) enable stable and efficient distributed BERT pre-training with batch size up to 64K/32K ( for sequence length 128/512 , i.e. , 8M/16M tokens per batch ) while maintaining the sample-wise convergence speed . For encoder-decoder models such as T5 , Raffel et al . ( 2020 ) use batch size up to 2K ( for sequence length 512 , i.e. , 1M tokens per batch ) . For autoregressive models such as the GPT family ( Radford et al. , 2018a ; b ; Brown et al. , 2020 ) , existing works use batch size up to 1.6K ( for sequence length 2K , i.e , 3.2M tokens per batch ) . Despite the benefit of reduced communication overhead , large-batch training is sensitive to hyperparameters and often leads to issues such as slow convergence , training instabilities , and model divergence ( You et al. , 2020 ; Li et al. , 2021 ) . Curriculum learning Inspired by how humans and animals are trained , curriculum learning aims to improve machine learning model training convergence speed by presenting easier ( or less complex ) examples earlier during training ( Elman , 1993 ; Sanger , 1994 ; Bengio et al. , 2009 ) . Previous studies have demonstrated the benefit of faster convergence speed by curriculum learning in many domains such as natural language processing , computer vision , and neural evolutionary computing . In this section we will discuss recent related works in the natural language processing area . For a more complete literature review , we recommend these recent survey papers ( Wang et al. , 2020 ; Soviany et al. , 2021 ) . In the natural language processing area , most of the curriculum learning works focus on small-scale one-stage tasks and downstream fine-tuning tasks , such as neural machine translation ( NMT ) ( Kocmi & Bojar , 2017 ; Bojar et al. , 2017 ; Zhang et al. , 2018 ; Platanios et al. , 2019 ; Zhang et al. , 2019 ) and natural language understanding ( NLU ) ( Sachan & Xing , 2016 ; 2018 ; Tay et al. , 2019 ; Xu et al. , 2020 ) . These works show that curriculum learning can improve convergence speed , reduce training time , and improve accuracy . In these works , the curriculum difficulty metrics for each training sample are usually defined as the sentence length , vocabulary frequency , the inference loss on smaller/cheaper models , or based on self-paced learning ( Kumar et al. , 2010 ) . For the pacing function ( i.e. , to decide the curriculum difficulty range when sampling next training data batch ) , these works usually use fixed predefined functions ( e.g. , gradually increase difficulty upper bound by linear , root , and exponential functions ) , bucketing heuristics ( group data with similar difficulties , and sample from a subset of buckets every time ) , or based on self-paced learning . Only a few works explore curriculum learning for language model pre-training . Press et al . ( 2020 ) apply curriculum learning to neural language modeling pre-training , specifically a transformer model with 247M parameters ( Baevski & Auli , 2018 ) . They find that by adding an additional first training stage with a shorter sequence length , it is possible to achieve the same dev . set perplexity with shorter total training time . Zhang et al . ( 2021 ) apply curriculum learning to BERT-base pre-training , a transformer model with 110M parameters . They find that by grouping sequences with similar length and with curriculum learning , it is possible to achieve similar downstream task accuracy with shorter pre-training time . Campos ( 2021 ) apply curriculum learning to ELMo pre-training , a bi-directional LSTM model with 93.6M parameters ( Peters et al. , 2018 ) . They test a variety of curricula on the WikiText-2 and WikiText-103 but do not find strong evidence that the use of curriculum learning can improve language model pre-training .

Summary: The paper examine optimization stabilities encountered during training of GPT-2 models. It demonstrates that as model size, learning rate, and batch sizes grow training dynamics of the model become unstable. To alleviate these issues, the authors propose a curriculum learning approach based on the length of the sequence. They show that this approach stabilizes the learning dynamics and drastically improves the training efficiency.