Efficient Stagewise Pretraining via Progressive Subnetworks
Abhishek Panigrahi Nikunj Saunshi Kaifeng Lyu Sobhan Miryoosefi Sashank Reddi Satyen Kale Sanjiv Kumar
Recent developments in large language models have sparked interest in
efficient pretraining methods. A recent effective paradigm is to perform
stage-wise training, where the size of the model is gradually increased over
the course of training (e.g. gradual stacking (Reddi et al., 2023)). While the
resource and wall-time savings are appealing, it has limitations, particularly
the inability to evaluate the full model during earlier stages, and degradation
in model quality due to smaller model capacity in the initial stages. In this
work, we propose an alternative framework, progressive subnetwork training,
that maintains the full model throughout training, but only trains subnetworks
within the model in each step. We focus on a simple instantiation of this
framework, Random Path Training (RaPTr) that only trains a sub-path of layers
in each step, progressively increasing the path lengths in stages. RaPTr
achieves better pre-training loss for BERT and UL2 language models while
requiring 20-33% fewer FLOPs compared to standard training, and is competitive
or better than other efficient training methods. Furthermore, RaPTr shows
better downstream performance on UL2, improving QA tasks and SuperGLUE by 1-5%
compared to standard training and stacking. Finally, we provide a theoretical
basis for RaPTr to justify (a) the increasing complexity of subnetworks in
stages, and (b) the stability in loss across stage transitions due to residual
connections and layer norm.