SimiGrad: Fine-Grained Adaptive Batching for Large Scale Training using Gradient Similarity Measurement

1 Introduction . Recent years have witnessed the success of deep learning in a variety of applications , such as computer vision [ 1 ] and natural language processing [ 2 , 3 ] . To push the performance boundaries , deep learning models in many applications have become increasingly complicated , and training such models requires large amounts of data as well as high cost . Taking GPT3 [ 3 ] ( an autoregressive language model that uses deep learning to produce human-like text ) as an example , the state-of-the-art model has 175 billion parameters and is trained on 300 billion data , which costs $ 12 million for a single training run [ 4 ] . Therefore , it is of critical importance to study scalable and efficient training for large scale machine learning [ 5 , 6 ] . The core idea of large scale machine learning is to parallelize model training across vast amounts of computing resources to speed up the training process . However , the overhead of the model synchronization due to massive parallelism in training can be the bottleneck for large scale training . To reduce the model synchronization overhead , large-batch training is a commonly employed strategy that can reduce the synchronization frequency . Taking BERT training as an example , increasing batch size from 256 to 4k can increase the training speed 2.5x faster on a 16-GPU cluster . OpenAI uses a batch size of 3.2 million to train GPT3 over a supercomputing cluster [ 3 ] , which accelerates the training speed by approximately over 500x compared to sample-by-sample run on a 16-GPU cluster . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Despite the benefits of large batch size , large batch size may degrade the model generalization performance [ 7 , 8 , 9 ] . To address this issue , existing works [ 10 , 11 , 12 ] demonstrate that a carefully tuned learning rate scheduling can help to mitigate such performance loss caused by large batch size . However , manually tuning learning rate requires significant amounts of efforts , which limits its potential use for wide adoption . [ 13 , 14 ] develop automated learning rate tuning methods for large batch size , which , however , rely on users to set appropriate batch sizes . [ 15 ] develops adaptive batching using second order information , aiming to reduce the generalization impact due to large batch sizes . [ 16 ] uses the gradient variance information to guide the batch size and learning rate configurations . Due to the high cost of analyzing the gradient variance and the second order information , existing methods can adapt the batch size only at a coarse-grained level ( i.e. , at the end of each epoch ) . It is worth noting that within each epoch , the variance of gradient may change swiftly and significantly [ 17 , 18 ] ( also see Section 4.1 ) . Therefore , the course-grained methods can not find a suitable batch size to reduce the performance loss , and thus miss the opportunity to achieve larger batch sizes with less performance impact . To address the above challenges , we propose a fine-grained batch size adaption approach SimiGrad to optimize the batch size at the iteration level . SimiGrad is driven by a novel gradient similarity measurement that can accurately capture the gradient variance information in a lightweight manner . We use the proposed gradient similarity measurement to derive the optimized batch size and learning rate with an automated algorithm . We integrate SimiGrad into mainstream machine learning frameworks by providing an open-sourced plugin tool . We evaluate SimiGrad via extensive experiments using popular benchmarks , such as CIFAR10 , ImageNet , and BERT-Large . The experimental results demonstrate that SimiGrad can consistently outperform state-of-the-art methods in both the average batch size and the model generalization performance . Particularly , we achieve a new state-of-the-art batch size of 78k in BERT-Large pretraining with the SQuAD score 90.69 , compared to 90.58 reported in the state-of-the-art LAMB [ 19 ] with 59k batch size . We also achieve better Pareto frontier for the trade-off between the batch size and the model performance , compared with state-of-the-art methods . We summarize our main contributions as follows : • We develop an accurate gradient similarity measurement for capturing the gradient variance at the iteration level . This measurement can be directly inferred from the gradients of replicas with low computation overhead . • We develop a fully automated fine-grained adaptive batching algorithm that can identify the optimal batch size and learning rate using the gradient similarity measurement . • We integrate SimiGrad into mainstream machine learning frameworks and open-source it 1 . • SimiGrad enables a record breaking large batch size of 77k for BERT-Large pretraining . With very large batch sizes , SimiGrad is able to achieve better performance than state-of-the-art methods . 2 Background and Related Work . 2.1 Mini-batching Training . The objective of training a neural network is to optimize the neural network parameters θ by minimizing the loss function L ( θ ) . One typical method of training a neural network is the gradient descent method [ 20 ] , which computes the gradient of L ( θ ) with respect to θ over the entire dataset ; however , it could be too expensive for training deep neural networks with large training datasets . A typical solution is based on mini-batch stochastic gradient descent ( SGD ) [ 21 ] , which samples a small subset , namely mini-batch , from the dataset and approximates the gradient using this batch instead . Such approximation introduces inevitable noise to the gradients . By increasing the batch size ( i.e. , the number of samples ) , we can reduce the noise level , but it is a diminishing return , meaning that after a sweet point we can gain almost no more information by further increasing the batch size [ 22 ] . In other words , the gradient computed on a too small batch is biased with noise , whereas a too large batch leads to a waste of time and insufficient updates to the model [ 23 ] . This makes the batch size a key control knob , and the choice of the batch size is of critical importance . In practice , a large batch may be too large to fit into GPU memory . The solution [ 24 , 25 ] is to break it into many micro batches and accumulates the gradients of all micro batches . This method is called gradient accumulation , which is widely used in large scale training . 2.2 Large Batch Size Training . For large scale neural network training , it is desirable to have a large batch size , as it improves hardware utilization and system scalability as well as reduce wall clock training time . However , there 1Code repo is available at : https : //github.com/HeyangQin/SimiGrad is a dispute as to how exact large batch size would impact the model performance . In [ 7 , 8 ] , it is reported that training neural network model with a large batch size leads to a non-trivial generalization gap compared with a small batch size . [ 10 ] suggests that this gap could be closed by adjusting the training regime . [ 11 ] observes from a large number of training runs that this gap could be mitigated if the learning rate is well-tuned for a large batch size . [ 12 ] further observes that the range of “ good ” learning rates is shrinking as the batch size grows . [ 9 ] further argues that with carefully tuned hyper parameters , the smaller batch sizes still outperform large ones in terms of test accuracy . Despite the disagreement over large batch size ’ s generalization performance , much effort has been made on how to choose an appropriate batch size and the corresponding learning rate for neural network training . [ 26 ] formulates SGD as stochastic differential equation and suggests that there exists an optimal noise scale related to the batch size , the learning rate and the dataset . Along this line , [ 16 ] formulates the noise scale in a different way and tests it in a variety of applications . Another line of research explores the potential of large batch size training from a practical perspective by fine tuning the training setting , especially the learning rate . [ 24 ] suggests that the learning rate should be increased proportionally to the square root of the increase ratio of the batch size . [ 13 ] proposes the warmup learning rate scheduling that can scale up the batch size to 8k for ResNet-50 on ImageNet . The most closely related works to this paper are AdaScale SGD [ 14 ] and gradient noise scale [ 16 ] . AdaScale SGD [ 14 ] proposes a novel learning rate scaling rule based on the gradient variance for training with a large batch size ; however , its performance hinges on the choice of a suitable batch size , which is challenging to determine in practice . Gradient noise scale [ 16 ] proposes to predict the batch size by comparing the gradients before and after the allreduce call in the data parallelism as an alternative implementation . SimiGrad differs from gradient noise scale [ 16 ] in three ways : 1 ) the approach in [ 16 ] highly depends on the parallelism degree of the system , whereas SimiGrad supports flexible parallelism degree ; 2 ) the measurement proposed in [ 16 ] is noisy and needs to use a hand-tuned exponential decay parameter to stabilize it ; and 3 ) [ 16 ] assumes the learning rate is well-tuned and close enough to optimal , whereas SimiGrad has no assumptions on learning rate . In particular , by evaluating [ 16 ] in a common cluster of 16 GPUs , we observe a very high noise level . We also observe that the noise level decreases as the number of GPUs increases , which indicates that the design in [ 16 ] highly depends on the parallelism degree . In comparison , SimiGrad can derive relatively stable measurements in a parameter free and system independent manner as shown in Fig . 3 . 2.3 Adaptive Batching . Using a fixed batch size throughout the training may not be the best practice . A desirable approach is to use adaptive batch sizes during training , which is suggested by research from both theoretical [ 22 ] and practical [ 27 , 28 ] perspectives . [ 22 ] converts the existing learning rate schedule into the batch size adjustment schedule . Yet the optimal learning rate is difficult to obtain in practice and the batch size adjusted in this way has limited scalability . [ 27 , 28 ] use fixed multi-step batch size schedulers that closely resemble the multi-step learning rate scheduler . Recent works focus on utilizing the gradient noise variance information to adjust batch sizes by various rules and measurements . [ 29 ] proposes a coupled criterion for adjusting both the batch size and the learning rate . [ 15 ] adjusts the batch size using the second order information and adversarial training . [ 16 ] proposes a gradient noise scale to adjust the batch size . 3 Motivation and Challenges . Due to the high cost of analyzing the gradient variance and the second order information , existing methods can adapt the batch size only at a coarse-grained level ( i.e. , at the end of each epoch ) . Our key observation is that the gradient variance may change swiftly across iterations within one epoch ( see Fig . 2 in [ 17 ] and Fig . 3 in [ 18 ] ) , which can not be captured by the existing methods . Taking [ 16 ] as an example , to get the gradient variance information , it needs to run at least 20 batches with different batch sizes . To perform such analysis at each iteration , the extra overhead due to gradient variance analysis would be more than 20 times of the original training , which is infeasible in practice . Due to these issues , the batch size usually needs to be compromised and large scale training methods that achieve state-of-the-art performance still heavily rely on manual tuning efforts [ 5 , 6 ] . To address these issues , one promising solution is by improving the granularity of gradient analysis , which , however , is challenging due to the following reasons . • To compute the gradient variance , we need to first collect the gradient information from batches at different sizes and then take large amounts of computing time and resources . For example , computing the gradient variance for a single 8k batch in BERT training takes 441ms and the time grows significantly with the batch size . • Even if we can obtain the gradient variance , it is not straightforward to derive the optimal batch size as the gradient covariance is a big matrix . • Existing batch adaption methods usually require a fine-tuned learning rate scheduler to achieve a good performance . Fine tuning the learning rate at the iteration level requires significant amounts of manual efforts , which may not be feasible in practice . • Due to the frequent gradient analysis , the system deployment is no longer a trivial engineering problem as the analysis and tuning overhead can disrupt the well optimized training pipeline in existing machine learning frameworks . For example in [ 16 ] , executing multiple batches with different batch sizes to compute the gradient variance is not supported by the training pipeline and may lead to inefficient memory use .

Summary: Large batch sizes are required to accelerate the large-scale training of deep learning (DL) models, which often comes at the cost of generalization performance. This paper introduces SimiGrad, an automated method that adapts the batch size at the mini-batch level during DL training. The proposed method achieves state-of-the-art performance on BERT-Large when employing very large batch sizes.   The first contribution of the paper is the use of the cosine similarity metric to measure the gradient variance of the gradient of a mini-batch. This replaces the typical expensive step that needs to calculate multiple batches and collect gradients from different batch sizes.   SimiGrad also automatically optimizes the learning rate of the training procedure. The automated adaptation of the learning rate used by SimiGrad is shown to out-perform other learning rate scaling rules when using large batch sizes.   SimiGrad is integrated into existing ML systems and the experimental results show that state-of-the-art results can be obtained when SimiGrad is applied to BERT-Large training using a record-breaking batch size of 77K.