Large Batch Optimization for Deep Learning: Training BERT in 76 minutes
Table 1: We use the F1 score on SQuAD-v1 as the accuracy metric. The baseline F1 score is the score obtained by the pre-trained model (Bert-Large) provided on Bert’s public repository (as of February 1st, 2019). We use TPUv3s in our experiments. We use the same setting as the baseline: the first 9/10 of the total epochs used a sequence length of 128 and the last 1/10 of the total epochs used a sequence length of 512. All the experiments run the same number of epochs. Dev set means the test data. It is worth noting that we can achieve better results by manually tuning the hyperparameters. The data in this table is collected from the untuned version.
['Solver', 'batch size', 'steps', 'F1 score on dev set', 'TPUs', 'Time']
We first discuss empirical results for speeding up Bert training. For this experiment, we use the same dataset as Devlin et al. The F1 score on SQuAD-v1 is used as the accuracy metric in our experiments. All our comparisons are with respect to the baseline Bert model by Devlin et al. To train Bert, Devlin et al. first train the model for 900k iterations using a sequence length of 128 and then switch to a sequence length of 512 for the last 100k iterations. This results in a training time of around 3 days on 16 TPUv3 chips. To ensure a fair comparison, we follow the same SQuAD fine-tune procedure of Devlin et al. As noted earlier, we could get even better results by changing the fine-tune configuration. For instance, by just slightly changing the learning rate in the fine-tune stage, we can obtain a higher F1 score of 91.688 for the batch size of 16K using Lamb. Below we describe two different training choices for training Bert and discuss the corresponding speedups. For the first choice, we maintain the same training procedure as the baseline except for changing the training optimizer to Lamb. We run with the same number of epochs as the baseline but with batch size scaled from 512 to 32K. The choice of 32K batch size (with sequence length 512) is mainly due to memory limits of TPU Pod. By using the Lamb optimizer, we are able to achieve a F1 score of 91.460 in 15625 iterations for a batch size of 32768 (14063 iterations for sequence length 128 and 1562 iterations for sequence length 512). With 32K batch size, we reduce Bert training time from 3 days to around 100 minutes. We achieved 49.1 times speedup by 64 times computational resources (76.7% efficiency). We consider the speedup is great because we use the synchronous data-parallelism. There is a communication overhead coming from transferring of the gradients over the interconnect. For ResNet-50, researchers are able to achieve 90% scaling efficiency because ResNet-50 has much fewer parameters (# parameters is equal to #gradients) than Bert (25 million versus 300 million).