Staged Training for Transformer Language Models

TL;DR

This work tackles the high compute cost of training large transformer language models by introducing staged training with growth operators that progressively enlarge a model from a small to a large one. The width and depth growth operators are designed to be loss-preserving and training-dynamics-preserving, and they operate on the full training state, including optimizer state and learning-rate schedules. A principled, schedule-based approach built on scaling laws (Kaplan et al.) yields compute-efficient growth timings, and practical methods estimate the necessary constants from small models. Empirically, the approach achieves up to 22% compute savings on GPT-2–style autoregressive models while matching the target model’s performance; the authors also provide public code and show that the method outperforms prior growth techniques that do not preserve loss or training dynamics. Overall, staged training with loss- and dynamics-preserving growth offers a practical pathway to more compute-efficient scaling of transformer language models, with broad applicability to autoregressive and possibly other transformer families.

Abstract

The current standard approach to scaling transformer language models trains each model size from a different random initialization. As an alternative, we consider a staged training setup that begins with a small model and incrementally increases the amount of compute used for training by applying a "growth operator" to increase the model depth and width. By initializing each stage with the output of the previous one, the training process effectively re-uses the compute from prior stages and becomes more efficient. Our growth operators each take as input the entire training state (including model parameters, optimizer state, learning rate schedule, etc.) and output a new training state from which training continues. We identify two important properties of these growth operators, namely that they preserve both the loss and the "training dynamics" after applying the operator. While the loss-preserving property has been discussed previously, to the best of our knowledge this work is the first to identify the importance of preserving the training dynamics (the rate of decrease of the loss during training). To find the optimal schedule for stages, we use the scaling laws from (Kaplan et al., 2020) to find a precise schedule that gives the most compute saving by starting a new stage when training efficiency starts decreasing. We empirically validate our growth operators and staged training for autoregressive language models, showing up to 22% compute savings compared to a strong baseline trained from scratch. Our code is available at https://github.com/allenai/staged-training.

Figures (7)

• Figure 1: We train a GPT2$_\textsc{large}$ (768M parameters) transformer language model by first training a model 1/4 the size (orange line), then increasing the model size by 4x by applying a growth operator to the entire training state, and restarting training (green line). The result is a large size model with comparable loss to one trained from scratch (blue line) but with reduced compute cost illustrated initially by the dashed red arrow.
• Figure 2: Our growth operators are loss-preserving and training-dynamics preserving. Using (a) as an example, GPT2$_\text{large/4}$ is the original model which is 4x smaller than the target model GPT2$_\textsc{large}$. The model GPT2$_\text{large/4x4}$ is the grown model resulting from growing GPT2$_\text{large/4}$ by 4x (doubling model width). The pre-growth point is highlighted on the original model GPT2$_\text{large/4}$, and the post-growth point is highlighted on the grown model GPT2$_\text{large/4x4}$. The pre-growth and post-growth points have the same loss, showing that the width growth operator is loss-preserving. To demonstrate that it is also training dynamics-preserving, we overlay the loss curve for the grown model over the target model and confirm the rate of loss decrease with respect to the number of tokens is the same as the target model trained from scratch. The x-axis is number of training tokens since random initialization, or from the start of the training stage (for GPT2$_\text{large/4x4}$). A similar result is seen in (b) for the depth growth operator.
• Figure 3: Comparing with three different baselines from prior work and ablation studies. $\bullet$ The width growth operator of gu2021transformer in GPT2${_\text{base/4x4width ffn}}$ and the depth growth operator of gong2019efficient in GPT2$_\text{base/2x2depth copy}$ are not loss preserving (higher initial loss). Also, GPT2${_\text{base/4x4width ffn}}$ is significantly underperforming the target model GPT2$_\textsc{base}$. $\bullet$ Resetting the optimizer state to zero instead of growing it (the zero_opt runs) have large instabilities and not preserving of the training dynamics. $\bullet$\ref{['subfig:grow_earlylate_tokens']} shows that restarting the learning rate schedule as in rae2021scaling is not training-dynamics-preserving. $\bullet$\ref{['subfig:grow_earlylate_tokens']}, \ref{['subfig:grow_earlylate_compute']} also show that not following our optimal schedule and grow the model too early or too late is still loss-preserving and training-dynamics-preserving but leads to lower compute saving
• Figure 4: Similar to Figure \ref{['fig:width_depth_token']}, our width and depth growth operators are loss-preserving and training dynamics preserving. $\bullet$ GPT2$_\text{large/16x16-width}$ indicates starting from a 16x smaller model then growing it 16x by doubling the width twice. $\bullet$ GPT2$_\text{large/16x4x4-width}$ indicates growing the model 16x over two stages, by doubling the width once, continue training the model, then doubling the width again. $\bullet$ The same applies to the depth growth operator.
• Figure 5: Growth operator and total compute. Our grown models are saving compute compared with target model.