Layer-Parallel Training for Transformers

TL;DR

The paper tackles the bottleneck of training very deep transformer models by introducing layer-parallel training via a neural ODE formulation and multigrid-in-time (MGRIT). By reinterpreting forward and backward passes as time stepping over layers and applying a coarsening factor $c_f$, the method exposes parallelism across the layer dimension while remaining compatible with data and model parallelism. A key contribution is an adaptive scheme that monitors gradient inexactness with a convergence factor and switches to serial gradients when needed to preserve convergence, enabling both speedups and accuracy on BERT, GPT-2, ViT, and MT pretraining tasks. The results show substantial parallel speedups with near-serial accuracy in many cases, and a practical switching mechanism to maintain reliable training dynamics, paving the way for scalable training of very deep foundation models.

Abstract

We present a new training methodology for transformers using a multilevel, layer-parallel approach. Through a neural ODE formulation of transformers, our application of a multilevel parallel-in-time algorithm for the forward and backpropagation phases of training achieves parallel acceleration over the layer dimension. This dramatically enhances parallel scalability as the network depth increases, which is particularly useful for increasingly large foundational models. However, achieving this introduces errors that cause systematic bias in the gradients, which in turn reduces convergence when closer to the minima. We develop an algorithm to detect this critical transition and either switch to serial training or systematically increase the accuracy of layer-parallel training. Results, including BERT, GPT2, ViT, and machine translation architectures, demonstrate parallel-acceleration as well as accuracy commensurate with serial pre-training while fine-tuning is unaffected.

Figures (12)

• Figure 1: Layer-parallel transformer. The ParallelNet contains a time grid hierarchy with the coarsening rate denoted by $c_f$. Experiments use a fine level time-step of $h=1$.
• Figure 2: Two-level MGRIT pseudocode (left), MGRIT diagram with $c_f = 2$, $L = 2$ on 2 devices (right).
• Figure 3: The long term training behavior using sequential versus layer-parallel with multiple GPUs. On the left, the validation accuracy for the MC example with 64 transformer layers, $L=2$, and $c_f=2$. On the right, the validation BLEU for the MT example with 6-6 transformer layers, $L=2$, and $c_f=3$. The plot corresponding to "2-->1" label illustrates a switch from parallel training with 2 GPUs to serial training with 1 GPU. Note that two depicted "2-->1" runs switch from a parallel to serial run at different points during the training.
• Figure 4: Plots of the loss for serial (blue), pure parallel (red) and switching to serial from parallel (green) for the BERT (left), GPT (middle) and ViT (right). In all the experiments, we see that purely layer-parallel runs will diverge from serial training after a certain point. However, one can recover the original dynamics by switching from parallel to serial at an appropriate time given by the indicator. The gray color in the BERT subplot indicates the min/max over three different seeds.
• Figure 5: The indicator values for BERT (forward in red, backward in blue), ViT and GPT (backward in blue) using MGRIT. We see that at the 70000th batch, 1000th, and 6000th batch respectively, the indicators exceed 1, meaning that one should switch to exact gradient computation then.