Progressively Stacking 2.0: A Multi-stage Layerwise Training Method for BERT Training Speedup
Authors: Cheng Yang, Shengnan Wang, Chao Yang, Yuechuan Li, Ru He, Jingqiao Zhang
Keywords: BERT, Training speedup, Multi-stage training, Natural language processing
Abstract: Pre-trained language models, such as BERT, have achieved significant accuracy gain in many natural language processing tasks. Despite its effectiveness, the huge number of parameters makes training a BERT model  computationally very challenging.  In this paper, we propose an efficient multi-stage layerwise training (MSLT) approach to reduce the training time of BERT.  We decompose the whole training process into several stages. The training is started from a small model with only a few encoder layers and we gradually increase the depth of the model by adding new encoder layers. At each stage, we only train the top  (near the output layer)  few encoder layers which are newly added. The parameters of the other layers which have been trained in the previous stages will not be updated in the current stage. In  BERT training, the backward calculation is much more time-consuming than the forward calculation, especially in the distributed training setting in which the backward calculation time further includes the communication time for gradient synchronization. In the proposed training strategy,  only  top few layers  participate  backward calculation, while most  layers  only participate forward calculation. Hence both  the computation  and communication efficiencies are greatly improved.  Experimental results show that the proposed method can greatly reduce the training time without significant performance degradation.