The Stability-Efficiency Dilemma: Investigating Sequence Length Warmup for Training GPT Models
Anonymous Author(s)
AffiliationAddress email
Abstract
Recent works have demonstrated great success in pre-training large-scale autoregressive language models (e.g., GPT-3) on massive GPUs.
To reduce the wall-clocktraining time, a common practice is to increase the batch size and learning rate.
However, such practice is often brittle and leads to a so-called stability-efficiencydilemma: increasing the batch sizes and learning rates leads to better trainingefficiency but can also result in training instability, leading to poor generalizationaccuracy or failed runs.
To better understand this phenomenon, we conduct anin-depth analysis on large-scale pre-training experiments replicating the GPT-2model with public dataset.
We find that there is a strong correlation betweentraining instability and extreme values of gradient variance.
We further identify thatsamples with long sequence lengths contribute to these extreme gradient variancevalues, especially at the beginning of the training, indicating that long sequencelength can be a main source of training instability.
Based on the analysis, we present a simple yet effective Sequence Length Warmupmethod that aims to solve the training stability-efficiency dilemma by avoidingextreme gradient variance values.
Moreover, we present a lightweight tuningstrategy that allows us to tune our method with just a small portion of the expensivefull training.
Experiments replicating GPT-2 models (117M and 1.5B) show thatour approach enables stable training with 8x larger batch size and 4x larger learningrate, whereas the baseline approach struggles with training instability.
To achievethe same or better zero-shot evaluation results, our method reduces the requirednumber of training tokens and wall clock time by up to 2.2x and 3.7x, respectively.
Experiments replicating GPT-3 model (125M) show that our approach enablesstable training with 8x larger batch size and 40x larger learning rate, and retains99% of the zero-shot accuracy on 11 tasks using 10x less data and 12x less timecompared to the original GPT-3 training recipe, while the baseline diverges underthe same settings and only retain 95% of accuracy under lower learning rate.
Introduction
Large-scale Transformer-based language models have achieved great success in many natural language processing tasks [44, 11].
Among them, large-scale autoregressive models, such as GPT-3 [6],have attracted lots of attention due to their superior performance on zero-shot generalization, i.e.,they can perform a wide range of tasks despite that they are not explicitly trained on.
However,pre-training GPT models raises huge challenges on training efficiency and less-discussed traininginstability issues.
On the efficiency side, as the model size continues to grow from a few hundreds ofmillions (e.g., GPT [29]), to billion-scale parameters (1.5B GPT-2 [31]), and to more recent hundredsof billions of parameters (175B GPT-3 [6]), the training cost also increases exponentially: it requires9.2 days on 512 V100 GPUs to train a 8.3B GPT-2 [38], and 47.8 days on 2240 A100 GPUs to train
a
530B GPT-3-style model [24].
Such a long training time makes it almost infeasible for most researchlabs and practitioners to reproduce these models.
Various solutions have been proposed to reduce thetraining wall clock time of these large models [38, 23, 33].
However, many solutions require usingmore GPUs or sophisticated system techniques.
In this work, we investigate speeding up the pre-training of GPT-style models via exploiting dataefficiency, not at the cost of excessive hardware resources.
In particular, in a distributed trainingenvironment, increasing the batch sizes and/or using more aggressive learning rates can make themodel converge faster [39].
However, it has been observed that larger batch sizes and learning ratescan make large-scale GPT model training more difficult, e.g., causing training instability that leadsto divergence or slow convergence [47, 9].
To investigate this training instability issue, we conducta thorough study of the GPT-2 pre-training task [31, 38] with different models sizes under variousbatch sizes, learning rates, and sequence lengths.
We find an stability-efficiency dilemma :
• A larger batch size (and larger learning rate) increases the per-iteration computational efficiency
but with increasing likelihood of training instability and even divergence.
• A smaller batch size makes the training more stable but decreases the per-step computationefficiency significantly.
We find it difficult to overcome this dilemma by existing techniques such as extra gradient clipping.
More recent proposed techniques such as batch size warmup proposed in [6] does not provide stabilitybenefit in our evaluations.
Recently, Shortformer[28] shows that by adding an additional first trainingstage with a shorter sequence length, language models can achieve the same dev.
set perplexity withshorter total training time.
However, (1) its main focus is to purely improve training efficiency insteadof solving the efficiency-stability dilemma, and our evaluations show that Shortformer’s 2-stageapproach is insufficient for overcoming the training instability issue for large models (Section 5.1),(2)
it is primarily evaluated on small scale transformer models (247M) on WikiText datasets (103Mtokens) without considering large-scale generative models like GPT with hundreds or even billions ofparameters, and (3) it does not discuss how to choose some of the important hyperparameters, whichare very expensive to figure out for large-scale model training.
Inspired by Shortformer, we investigate the importance of sequence length in training GPT models andfind that it plays an important role in both training stability and efficiency.
Based on our investigation,we propose a simple yet effective method called Sequence Length Warmup (SLW) , which startstraining with short sequences and gradually increases the length.
We observe that our approachenables stable and efficient training with much larger batch sizes and learning rates than baselineapproaches.
Specifically, we make the following contributions: (1) We conduct an extensive study ofthe GPT-2 pre-training task, which provides detailed insights about the training stability-efficiencydilemma, the correlation between instability and gradient variance outliers, and how sequence lengthplays a critical role (Section 3).
(2) Based on the study, we present a simple yet effective sequencelength warmup method for GPT-style model (and autoregressive model in general) that enables stabletraining with improved training efficiency.
We also identify a lightweight hyperparameter tuningstrategy for the approach, which identifies promising hyperparameters by only incurring a smallfraction of the expensive total pre-training cost (Section 4).
The implementation of our approach aswell as the necessary changes to the GPT-2/3 pre-training framework has been open sourced in adeep learning optimization library (name hidden to maintain anonymity).
(3) We conduct large-scaleexperiments to demonstrate the proposed work’s ability to provide superior training stability andefficiency at the same time (Section 5).
Our empirical results show that:
• SLW enables stable and efficient training with 8x larger batch size and
4x larger learning rateon GPT-2 (117M and 1.5B) and GPT-3 (125M) models with public datasets, while the baselineand relate works struggle with instability under the same settings.
To achieve the same or betterzero-shot WikiText-103/LAMBADA evaluation results at the end of training, SLW reduces therequired number of training tokens and wall clock time by up to 2.2x and 3.7x, respectively.•
For GPT-3 experiments we study an even more aggressive training scenario where only 10% of datacan be used.
Our method, with 8x larger batch size and 40x larger learning rate than the originalGPT-3 training recipe, is able to maintain the training stability, retaining 99% of the zero-shotaccuracy on 11 evaluation tasks, and use 10x less data and 12x less time.
Without our method, thebaseline has unrecoverable divergence under the same settings, and can only retain 95% of thezero-shot accuracy after lowering learning rate to 30x.
Background and Related Work
Language Model Pre-training: The accuracy of transformer-based language models grows substantially with its model size [29, 31, 6].
Today, a large language model such as GPT-3 [6] contains up to175B parameters, and recent studies show that model accuracy can continue to improve with evenlarger model sizes [16].
However, training these large models often incurs excessively long trainingtime and training difficulties [6].
Therefore, there are a lot of demands of performing efficient andstable 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, thespeedup gains often saturate beyond a few tens of GPUs, because communication becomes the majorbottleneck, i.e., the workers will spend more time communicating gradients than computing them, asthe number of GPUs increases.
To mitigate this bottleneck, recent works such as 1-bit Adam [41] havestudied gradient compression and demonstrate their effectiveness against auto-encoding models suchas BERT [11].
An alternative approach to alleviate these overheads is to use large batch sizes.
Forexample, LAMB [50] and 1-bit LAMB [20] enable stable and efficient distributed BERT pre-trainingwith batch size up to 64K/32K (for sequence length 128/512, i.e., 8M/16M tokens per batch) whilemaintaining the sample-wise convergence speed.
For encoder-decoder models, T5 [32] uses batchsize up to 2K (for sequence length 512, i.e., 1M tokens per batch).
For autoregressive models suchas the GPT family [29, 31, 6], 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-batchtraining is sensitive to hyperparameters and often leads to issues such as slow convergence, traininginstabilities, and model divergence.
For example, recently a research project shared that they aredealing with challenging training instability issues when pre-training a 104B GPT-style model withbatch size 2K [47], and another work on a 540B model with batch size 2K observed spikes in the lossroughly 20 times during training, despite the fact that gradient clipping was enabled [9].
Curriculum Learning: Our method can be viewed as a kind of curriculum learning (CL)
[12, 37, 2],which presents easier/simpler examples earlier during training and gradually increases the sampledifficulties 1 .
Comparing with traditional CL works which focus on solely improving the convergencespeed under the same batch size, learning rate and other hyperparameters, our work is motivatedby the stability-efficiency dilemma
and we aim to achieve both efficient convergence and betterstability by enabling stable training with more aggressive hyperparameters.
To our knowledge, weare the first to investigate and confirm that certain curriculum learning method can provide a dualstability-efficiency benefit.
Moreover, although CL was explored and verified for NLP one-stage andfine-tuning tasks [18, 5, 54, 27, 55, 34, 35, 42, 48], there are only a few works explore curriculumlearning for language model pre-training [28, 53, 7].
GPT-2 Pre-training Stability-Efficiency Analysis
In this section we perform an in-depth analysis of pre-training tasks (without our method) replicatingthe GPT-2 models with public data.
We follow the training pipeline from the NVIDIA Megatron-LMwork [38] 2 .
All of the experiments are performed on 128 NVIDIA V100 GPUs (32GB memory).
There are 16 nodes and 8 GPUs per node.
GPUs inside the same node are connected by NVLink2.0, and nodes are connected by a 100 Gigabit InfiniBand EDR inter-node network.
We evaluatetwo GPT-2 model sizes from the original GPT-2 work [31]: 117M parameters (12 layers, 768 hiddensize, 12 attention heads) and 1.5B parameters (48 layers, 1600 hidden size, 25 attention heads).
Fortraining data, we collect and use the same dataset blend as the Megatron-LM work: Wikipedia [11],
CC-Stories [43], RealNews [52], and OpenWebtext [30].
We evaluate two sets of training parameters.
The first set follows the Megatron-LM work: batch size512, 300K total training steps (157B tokens), and learning rate 1 .
5 × 10 − 4 with a linear warmup of3K steps and a single cycle cosine decay over the remaining 297K steps ( 1 × 10 − 5 min.
learning rate).
The second parameter set tests a more aggressive training strategy: batch size 4K ( 8 × larger), 37.5Ktotal training steps (157B tokens 3 ), and learning rate 6 × 10 − 4 ( 4 × larger) with a linear warmup of3K steps and a single cycle cosine decay over the remaining 34.5K steps (same min.
learning rate).
For sequence length/context size, we mainly use 1K which is the default for GPT-2.
But we alsotest 2K (on the 117M model with batch size 512 and 157B tokens) which is the default for GPT-3.
[Figure]
All experiments are performed with mixed precision/
FP16 training, Adam optimizer ( β 1 = 0 .
9 ,β 2 = 0 .
999 , ϵ = 1 × 10 − 8 ) [17], 0.01 weight decay, same random seed, and gradient clipping at 1.0.
[Table]
[Figure]
[Table]
The stability-efficiency dilemma:
Figure 1(a)and 1(b)
present the training loss curves of 5 baseline cases under different model sizes, batch sizes(and LR), and sequence lengths.
At 117M, the baseline has a few training loss spikes at batch size 4K. At1.5B, the baseline has many loss spikes when trainingwith either batch size 512 or 4K. As an indicative measurement to quantitatively study training instability,we define “loss ratio” which measures the ratio between the current step training loss and the minimumloss among all previous steps.
A ratio larger than 1.0means that current step’s loss is larger than the previous minimum loss, thus larger ratio indicates a largerloss spike and training instability.
Table 1 summarizesthe number of steps with loss ratio larger than 1.2, andthe maximum loss ratio during the training.
At 117Mmodel size only the baseline with batch size 4K hashigh loss ratios up to 1.421.
At 1.5B model size thebaseline with both batch size 512 and 4K has muchmore steps with large loss ratios, and with the maximum loss ratio as high as 5.65.
Baseline with batchsize 4K is less stable than baseline with batch size512, indicating that larger batch sizes could lead to more training instability risks.
In Appendix A.2.1we show that larger learning rates under the same batch size could also increase training instability.
Training instability are undesirable because (1) it could lead to divergence that never recover asin [47] and our GPT-3 experiments (Section 5.2); (2) in our GPT-2 case it leads to worse convergence,validation loss, and zero-shot downstream task accuracy.
Table 2 summarizes the zero-shot WikiText103/LAMBADA evaluation results.
For both 117M and 1.5B models, increasing baseline’s batchsize (and LR) or sequence length leads to worse evaluation results due to the training instability(case 1, 4, 7, 10, 13 in table).
On the other hand, increasing batch size (and LR) or sequence lengthimproves training efficiency, reducing the training time by up to 2.3x under the same number oftraining tokens (case 1, 4, 10, 13).
Overall, this demonstrates the stability-efficiency dilemma forbaseline pre-training: the training is more stable and can achieve better final generalization, butpresumably with poorer training efficiency under smaller batch size/learning rate/sequence length;increasing them leads to better training efficiency, but with lower stability and worse generalization.
The correlation between instability and gradient variance outliers: For stochastic gradientoptimization, when the gradient variance is large, the algorithm might spend much time bouncingaround, leading to slower convergence and potential divergence [46].
Previous studies show that
[Table]
variance reduction methods improve training stability in areas such as reinforcement learning [22, 8, 1]
[Table]
Figure 1(c), 1(d), 1(e) and 1(f) plot the l 1 norm and max element of Adam’s variance state ( √ v t ,where v t =
β 2 v t − 1 + (1 − β 2 )( g t ) 2 ) 4 .
When baseline’s batch size increases, the variance normdecreases but the max element increases.
Comparing GPT-2 117M and 1.5B cases, larger model sizeleads to larger variance norm and max element.
When sequence length increases for the GPT-2 117Mcase, the variance norm stays the same but the max element increases.
To further study the link between instability and gradient variance, Figure 1(g) and 1(h) plot the lossratio (defined earlier in this section) and gradient variance norm/max element (all normalized by maxvalue) for the most unstable 1.5B baseline with 4K batch size.
Results show that when training lossspike happens and loss ratio increases, the gradient variance norm/max also increase (especially themax outliers).
Table 3 presents the Pearson correlation coefficient calculations, which demonstratea statistically significant positive correlation between loss ratio and gradient variance norm/max.
Overall, our analysis shows that training instability has a strong correlation with gradient variancenorm and (especially) max element outliers.
Length of early data sequences is critical to training stability: Aiming to solve the stabilityefficiency dilemma we first tried traditional methods such as increasing gradient clipping, but itdoes not fully resolve the instability issue (Appendix A.2.2).
Seeing that in Figure 1 the traininginstability mostly happens at the first half of training, we then explored whether we can solve theissue by gradually increasing any of the batch size/learning rate/sequence length during training.
Wealready employed the same learning rate warmup mechanism used by existing GPT-2 and GPT-3works [31, 38, 6].
We tried the batch size warmup method proposed in GPT-3 work [6], but theinstability issue still appears when increasing the batch size (Section 5.1).
Our investigation on thesequence length leads to interesting findings, where we find that sequence lengths play a critical rolein training instability.
Figure 2 presents the training loss curve of the most unstable GPT-2 1.5Bpre-training with batch size 4K and seqlen 1K, together with another two artificial settings: one withseqlen 128, the other with mixed seqlen where we feed 900 steps of seqlen 128 then 100 steps ofseqlen 1K in every 1K steps.
The seqlen 128 case has no instability issue, even with large modelsize/batch size/learning rate.
The mixed seqlen case has instability issues, and (1) they mostly happenwhen we switch to seqlen 1K (e.g., at step 900, 1900, 2900...); (2) they mostly happen during the first5K steps, and after that it becomes more stable than the seqlen 1K case.
These observations indicatethat training instability is strongly correlated with early long sequence lengths, which motivates usto explore the sequence length warmup method described in the next section, and evaluations in
Section 5 will demonstrate how this method provides a gradient variance reduction effect and solvesthe stability-efficiency dilemma.
[Figure]
The Sequence Length Warmup Method
The analysis in last section about training instability and sequence lengths motivates us to exploresequence length warmup methods: the model needs to start learning with short sequence length formore stable training, then gradually increase the length when training becomes more stable so thatthe model can still learn from longer contextual information to achieve better final model accuracy.
The sequence length warmup strategy depends on two factors: how to support variable sequencelength during training and how to adaptively decide the sequence length for each iteration (the pacingfunction).
For the first component, we develop an efficient truncation-based implementation:
Forthe baseline GPT-2 pre-training, the raw text inputs are indexed into sequences with the same lengthbefore training, so that the model can efficiently retrieve a batch of fixed-length sequences regardlessof the actual sentence boundaries.
It’s possible to sort input data based on actual sequence lengths,but indexing the text based on all possible sequence lengths adds significant amount of overheaddue to the massive pre-training data.
To avoid the large indexing overhead, we take a lightweightapproach: our implementation still lets the dataloader index the raw text into only the full sequencelength.
At each training step, our method uses pacing function to determine the sequence length andthen truncates the full-length sequences to obtain a modified version of the mini-batch for training.
We define the pacing function as a step-wise linear function with the following properties: Givena starting sequence length seqlen s , an ending sequence length seqlen e (full sequence length), anda total duration T (number of steps), the sequence length used for the training batch at step t isseqlen t = seqlen s + ( seqlen e − seqlen s ) × min ( tT , 1) .
Besides step-wise linear, we also explored3 other pacing functions: i)
We tried a discrete 2-stage pacing function from [28], but it leadsto unstable training and worse convergence (Section 5.1).
ii)
We tried a step-wise root function( seqlen t = seqlen s + ( seqlen e − seqlen s ) × min (( tT ) r , 1) , where r is the root degree), whichperforms similar to linear but requires one extra hyperparameter.
iii)
We tried an adaptive pacingfunction based on training/validation losses, which also performs similar and requires extra tuning.
[Figure]
“SLW
Pacing function analysis and tuning strategy:
To study the impact ofour approach’s pacing function, we setthe starting sequence length ( seqlen s )fixed at 8 and perform a grid search forthe pacing function duration ( T ) on the
GPT-2 117M case full training (detailsin Appendix A.3).
We then choose the“best” duration based on the test dataperplexity and zero-shot evaluation results after full training.
All the caseshave quite comparable evaluation results, indicating that the performanceis not very sensitive to the duration Twithin a reasonable range.
This grid search sheds light on a low-cost tuning strategy: we find thatthe best duration T is always the longest duration that does not have significant validation perplexityfluctuation during the first 10K steps (i.e., a few multiples of the LR warmup steps).
In Figure 3(a)the SLW 60K is the longest duration we tested that does not have significant validation fluctuationduring the first 10K steps.
In Figure 3(b) and Appendix A.3 SLW 60K does provide the best finalvalidation perplexity, best final test perplexity, and second best eval results.
Since it does not requiretraining the model until full convergence, this heuristic greatly reduces the hyperparameter tuningcost of our approach.
Another grid search on the starting sequence length seqlen s shows that it’s
[Figure]
generally better to set it as small as possible, to maximize the stability and convergence speedupbenefit.
However, if the validation perplexity has significant fluctuation near the starting sequencelength, increasing seqlen s would lead to better convergence.
Overall, the low-cost tuning strategy can be summarized as: (1) Start with seqlen s = 8 and T = afew multiples of LR warmup steps.
(2) Increase seqlen s until the validation perplexity no longerhas significant fluctuation at the very beginning.
(3) Perform a binary search to find the largest Tthat does not have significant validation perplexity fluctuation during the first few multiples of LRwarmup steps.
For the GPT-2 1.5B and GPT-3 125M models, we used this strategy to tune T andseqlen s for the pacing function, and results show that this low-cost tuning strategy could providesimilar stability-efficiency benefit as grid search on full training runs (GPT-2 117M case).
Evaluation5.1 GPT-2 experiments
For GPT-2 model, dataset, and hardware, we follow the same methodology in Section 3.
For proposedwork’s pacing function configurations (Section 4), we use seqlen s = 8 / 64 (for 117M/1.5B modelbased on tuning) and seqlen e = 1 K/ 2 K (full sequence length).
To fully utilize the NVIDIA Tensor
Core acceleration, we add a seqlen t = seqlen t − ( seqlen t mod 8) postprocessing to make sure thesequence length is always a multiple of 8.
For the total duration T , we tune this parameter (gridsearch for 117M and low-cost tuning for 1.5B) for each case.
For the training parameters, for ourapproach we use the same shared parameters as the baseline except two parameters: 1) Becauseduring sequence length warmup the number of tokens in a data batch is smaller, we modify thetraining termination condition so that all cases stop when reaching the same 157B training tokens. 2)
Because of 1), proposed approach now has more training steps, which make it necessary to modifythe learning rate decay schedule to have a fair comparison with the baseline.
We change the learningrate decay to token-wise over the 157B tokens (still cosine decay) instead of step-wise over the totalnumber of steps.
We describe the underlying rationale in Appendix A.1.
Based on the following observations, we demonstrate that our approach resolves the dilemma andsimultaneously improves the stability and efficiency.
We will mainly present the GPT-2 1.5B resultsand leave some GPT-2 117M results in Appendix.
Significant stability gain: In Section 3 Table 1 we discussed how we measure the training instabilitybased on the “loss ratio” metric, which shows that the baseline becomes less stable under largermodel size/batch size/learning rate/sequence length.
Comparing baseline and proposed work in thistable shows that our work reduces this instability measurement to zero in all cases, together with maxratio close to 1.0 (no spike).
This demonstrates the significant stability gain by our method.
Faster token-wise and time-wise convergence: Figure 4(a) and 4(b) present the validation perplexitycurves during GPT-2 1.5B seqlen 1K pre-training, comparing baseline and our approach.
When thebatch size increases from 512 to 4K for baseline, the time-wise convergence becomes faster but the token-wise convergence becomes slower and poorer.
On the other hand, our approach at batch size4K provides faster and better convergence both token-wise and time-wise comparing with the bestbaseline curve in each case.
Our approach with batch size 512 provides smaller convergence speedupbecause (1) Baseline with batch size 512 has less instability issue, limiting the gain from the proposedapproach; (2) At batch size 512 the communication overhead is very high, and our approach takesmore steps (i.e., communication rounds) than baseline to reach the same 157B training tokens.
Thisextra communication cost “cancelled” part of the time saving from our approach.
For GPT-2 117M,our approach provides similar token-wise and time-wise convergence speedup (Appendix A.3).
Advancing cost-quality Pareto curve: In Section 3 Table 2 we discussed about baseline’s zero-shotevaluation results.
For proposed work eval results in this table, we present them in two ways: oneevaluated at the earliest checkpoint that provides better eval results than baseline (batch size 512 andseqlen 1K); the other one evaluated at the end of full training.
Results show that our approach is ableto advance the cost-quality Pareto curve in two ways: (1) To reach the same eval result quality asbaseline, our approach reduces the required number of pre-training tokens and wall clock time by upto 2.2x and 3.7x, respectively; (2) Under the same 157B training tokens, our approach can furtherimprove the eval result quality.
In (1) the time-wise saving is higher than the token-wise because(a)
For each Transformer block, the self-attention and intermediate layers have time complexity of
O ( B × L 2 × H ) and O ( B × L × H 2 ) , respectively 5 .
The proposed method uses shorter sequences atthe beginning, reducing the time complexity quadratically for the self-attention sub-layer and linearlyfor the intermediate sub-layer of Transformer blocks; (b) By enabling stable training at larger batchsize, our approach achieves additional time-wise saving by reducing the communication overhead.
Variance reduction helps stabilize training:
In Section 3 we discussed the strong correlationbetween training instability and gradient variance norm/max.
Figure 4(c) and 4(d) demonstrate thatproposed approach stabilizes training and reduces both the Adam variance norm and the variancemax element.
Importantly, it avoids all the spikes of the variance max element, which all happen tobe where the baseline training diverges.
The shape of variance norm/max element shows that at theearly stage of training when the amount of training data is much smaller than the model capacity, themodel can easily overfit on the initial training data, causing extreme gradient variance and traininginstability.
Our approach, beyond traditional curriculum learning, acts as a regularization method andreduces the overfitting at the early stage of training, which is also why it has slower convergence infirst half of Figure 4(a).
One may wonder why gradient clipping cannot help avoid these extremegradient variance outliers.
Although gradient clipping can avoid large gradient at every single step, itcannot avoid the gradient variance getting accumulated from multiple steps (Appendix A.2.2).
Comparing with related works: We now compare the proposed work with two related works on themost challenging “1.5B model + batch size 4K” case.
The first work is the Shortformer where thefirst stage uses shorter sequences and the second stage uses full-length sequences [28].
Followingthe grid search in the paper, we use seqlen 128 for the first stage and set its duration at about halfof the baseline duration (20K steps).
The second work is the batch size warmup technique used by
GPT-3 [6], where we set the starting batch size at 128 and then gradually increase it to 4K, and setthe warmup duration same as the proposed work.
Other training hyperparameters are unchanged.
Figure 4(e) to 4(h) present the results.
Both related works provide convergence speedup but it is lessthan our work.
More importantly, they still have training instability issues.
The Shortformer has anobvious training divergence at step 20K when the sequence length switches from 128 to 1K (the spikeat 20K in Figure 4(h)).
This shows when staying at the same shorter sequence length for too long, themodel becomes heavily overfitted for that length which leads to divergence risk when/after switchingto full length.
Although both batch size warmup and our method reduce the number of tokens perbatch in a similar fashion, batch size warmup does not provide any training stability benefit comparedto the baseline.
This indicates that providing the same number of shorter (simpler) sequences leadsto better training stability than providing fewer number of same length (same difficulty) sequences.
In addition, batch size warmup has a limitation that the batch size must be multiple of data-parallelsize, which will be large for distributed training.
On the other hand, for our method the sequencelength only needs to be multiple of 8 to enable Tensor Core acceleration.
Both related works providenon-zero “loss ratio” in Table 1 and worse zero-shot evaluation results in Table 2.
5.2 GPT-3 experiments
For experiments replicating the GPT-3 125M model [6] using the Pile public dataset [13], first wereproduce the original GPT-3 training recipe: 300B training tokens, seqlen 2K, batch size 256 with5 B, L, H represent batch size, sequence length, hidden size. by original GPT-3 work [6] task eval results
[Figure]
[Table]
[Figure] [Figure]
batch size warmup (start with 16 then gradually increase to 256 in first 4B tokens), learning rate6 × 10 − 4 with a linear warmup of 375M tokens and a single cycle cosine decay over 260B tokens( 6 × 10 − 5 min.
learning rate) 6 .
Then we explore an aggressive training scenario where only 30Btokens (10%) are allowed.
This is because (1) GPT-3 paper admits that it has poor training sampleefficiency and it sees much more text during pre-training than a human sees in the their lifetime [6, 21].(2) There could exist cases where the total amount of data/computation resource is limited.
We adjustseveral hyperparameters in this 30B-token training: 8x batch size (2K) for better training efficiency,learning rate decay reduced to 30B tokens (based on study that LR schedule should match totaltraining tokens [14], warmup stays at 375M), min.
learning rate reduced to 0 (based on recent studyon GPT-3 [49]).
For baseline we keep using 4B-token batch size warmup, but when our method isused ( seqlen s = 72 , T = 11 .
5 K based on tuning)
we disable it since both methods reduce tokensper batch.
And for both cases we tune the learning rate and use the highest one that provides stabletraining, which is 30x ( 1 .
× 10 − 2 ) for baseline and 40x ( 2 . 4 × 10 − 2 ) for our method.
Figure 5 and 6 present the training loss and gradient variance max for the GPT-3 pre-trainingexperiments.
When applying 40x learning rate to the baseline (batch size warmup), it quicklydiverges and cannot continue to train due to NaN losses.
The corresponding gradient variance maxelement becomes a flat line after divergence because the gradients on all dimensions are so largethat all gradients get clipped including the max element.
After lowering the learning rate to 30x, thebaseline is able to finish the whole training, but it can only retain 95% of average zero-shot accuracy on11 tasks (HellaSwag [51], LAMBADA [26], TriviaQA [15], WebQs [3], Winogrande [36], PIQA [4],
ARC Challenge/Easy
[10], ANLI R1/R2/R3 [25])
compared with the case that reproduces the original
GPT-3 training recipe 7 as shown in Table 4.
In contrast, our approach enables stable training with40x learning rate, demonstrates lower gradient variance max outliers than baseline with 30x learningrate, retains 99% of the original training recipe’s average zero-shot accuracy, and achieves 10xdata saving and 12x time saving.
This demonstrates that the proposed method not only solves thestability-efficiency dilemma, but also opens a promising direction of significantly reducing totaltraining cost in a different data efficiency dimension.
Conclusion
This paper presents the Sequence Length Warmup method, which solves a stability-efficiency dilemmainside GPT-style model pre-training, an critical issue that hinders efficient language model pre-trainingas explained in our in-depth analysis.
By enabling stable training on more aggressive training recipe,this method also motivates a new dimension of training cost reduction by improving the data efficiency,as demonstrated by the 10x data and 12x time saving in our GPT-3 experiments.
We believe thatthe effectiveness, simplicity, and easy-to-use/tune make the proposed method a must-try for deeplearning practitioners, and we hope this work could motivate more studies on improving training dataefficiency.