STEP: Staged Parameter-Efficient Pre-training for Large Language Models (2024)

Abstract: Pre-training large language models faces significant memory challenges due to the large size of model weights. We propose STaged parameter-Efficient Pre-training (STEP), which combines ideas from parameter-efficient tuning and staged training. We conduct experiments on pre-training models of various sizes and demonstrate that STEP can achieve up to a 40.4% reduction in maximum memory requirement compared to vanilla pre-training while maintaining comparable performance.

heading: 1 Introduction
text: Large Language Models (LLMs) have become a fundamental technology in artificial intelligence. One challenge we aim to address in the research on LLMs is the vast amount of computational resources needed for pre-training, e.g., LLaMA (Touvron et al., 2023). This requirement for enormous computational resources is a significant obstacle to the research of LLMs.
To tackle this challenge, methods for reducing computational costs during pre-training have been actively studied. For example, ReLoRA (Lialin et al., 2024) reduces the computational cost by repeatedly applying low-rank adaptations while freezing the original parameters during pre-training. However, ReLoRA often degrades performance compared to vanilla pre-training under fair conditions (Lialin et al., 2024; Zhao et al., 2024); there is still considerable room to improve in this line of studies. From this background, this paper attempts to develop a method for pre-training LLMs that can achieve comparable performance at the same computational cost as vanilla pre-training while reducing the maximum memory requirements.
For this goal, we propose a method that combines ideas of Parameter-Efficient Tuning (PET) (He et al., 2022) and staged training (Shen et al., 2022). The basic concept is that by incorporating the idea of staged training, we can reduce the
maximum memory requirement by (1) pre-training a model with a smaller size in the first stage and (2) freezing the parameters already pre-trained in the previous stages and instead introducing much smaller additional training parameters following the PET technique in the remaining stages. Hereafter, we refer to our method as STaged parameter Efficient Pre-training (STEP). Figure 1 illustrates this concept.
We explore the effectiveness of STEP in pretraining experiments by comparing the baseline (Vannila pre-training) and conventional method (ReLoRA) under the same computational cost. We demonstrate that STEP achieves up to a 40.4% reduction in maximum memory requirements compared to vanilla pre-training while maintaining comparable validation es.

heading: 2 Related Work
text: Memory Efficient Training for LLMs Several memory-efficient training approaches have been actively developed in the literature of training LLMs (Rajbhandari et al., 2020; Korthikanti et al.,
607
2023). Among these, one major approach is training parameter reduction methods. One approach is PET (He et al., 2022), such as Adapter (Houlsby et al., 2019) and LoRA (Hu et al., 2022). PET techniques have mostly been developed for fine-tuning LLMs. There are only a few studies applying the PET techniques to the LLM pre-training (Lialin et al., 2024; Zhao et al., 2024). ReLoRA is a representative method designed for pre-training LLMs using LoRA (Hu et al., 2022). However, to achieve comparable performance to vanilla pre-training, ReLoRA needs to train the model in the vanilla setting for the first several steps. This implies that ReRoLA is currently unable to reduce the maximum memory requirement, as it requires the same amount of memory as vanilla pre-training.
Staged Training (Shen et al., 2022) The core idea behind Staged Training is based on the observation that while small-scale models are advantageous in the initial stages of learning from a computational efficiency perspective, large-scale models eventually achieve lower (Kaplan et al., 2020). Staged Training leverages this observation by training a small-scale model with high computational efficiency and applying an expansion operation called the Growth Operator during training. This operation expands the dimensions of Transformer layers and adds new layers. Regarding memory usage in staged training, since most existing studies train the full parameters of the model, this approach does not reduce the maximum memory requirements.

heading: 3 STEP: Staged Efficient Parameter Training
text: As briefly described in Section 1, our goal is to develop a method for pre-training LLMs that can achieve comparable performance at the same computational cost while reducing the maximum memory requirements during pre-training.

heading: 3.1 Procedure
text: The following four procedures are an overview of STEP and how it efficiently trains LLMs;
(Procedure 1) STEP performs a vanilla pretraining on a model with a much smaller size than the target model size as an initial model.
(Procedure 2) STEP expands the layers of the initial model to increase its size.
(Procedure 3) STEP also introduces the PET parameters given by the parameter-efficient adaptors for the layers trained in Procedure 1.
(Procedure 4) STEP continues to pre-train the parameters in layers newly added in Procedure 2 and the adaptors added in Procedure 3 while freezing those in layers trained in Procedure 1.
Note that The first to fourth red right-arrow in Figure 1 corresponds to Procedures 1 to 4, respectively. After finishing Procedure 4, we obtain the pre-trained model.1

heading: 3.2 Growth Layer Operator
text: This section explains how we expand the layers in Procedure 2. Given a model with n layers, the Growth Layer Operator modifies the structure of the model’s layers. We use Interpolation (Chang et al., 2018; Dong et al., 2020; Li et al., 2022), which adds new layers between existing layers and initialize them with the lower layer weights, namely ϕnew2i = ϕ new 2i−1 = ϕi.
We further extend it by incorporating an idea of a fusing method that averages the parameters of the two layers (O’Neill et al., 2021), namely, ϕnew2i = (ϕi+ϕi+1)/2, which we call InterpolationMean. The validity of using the average will be verified through experiments（Section 4.4). We discuss more detailed initialization in Appendix A.

heading: 3.3 Incorporating PET parameters
text: This section provides additional information about Procedure 3, which introduces PET parameters by the adaptors. We specifically focus on the low-rank adaptation method (Hu et al., 2022; Lialin et al., 2024) for this part.

heading: 3.4 Maximum memory requirement of STEP
text: We assume that the maximum memory requirement during the pre-training can be estimated by the size of model states, which include the parameters of the model itself, the gradients of the model parameters being trained, and the optimizer state.2 Moreover, we assume that we use a typical Transformer model and the Adam optimizer (Kingma and Ba, 2014), which are a commonly used configuration for pre-training LLMs. Additionally, we assume that all parameters are represented as 32-bit
1Note that we have the option to continue growing the layers by repeating Procedures 2 to 4.
2Other memory usages, such as activations, can be reduced using methods like Activation Recomputation (Korthikanti et al., 2023).
floating-point numbers. Consequently, when the number of parameters in one layer of the Transformer is Player and the number of layers in the model is n, the memory usage of the model state, expressed in bytes, is given by
Ptrn = 4n(Player︸ ︷︷ ︸ model +Player︸ ︷︷ ︸ gradient +Player+Player︸ ︷︷ ︸ optimizer )
= 16nPlayer,
(1)
where the optimizer state of Adam consists of two parts; the gradient momentum and variance.
Regarding the maximum memory requirement for STEP, let ni be the number of layers increased in the i-th stage from the i − 1 stage in STEP, where n0 = 0. Let Ni represent the total number of layers in the i-th stage model, namely, Ni =∑i
k=1 nk. Moreover, E(Player) denotes the number of parameters for the single layer, Player, added by PET. Then, we can estimate the maximum memory requirement for the stage i, that is, Mi, as follows:
P STEPi =    16niPlayer if i = 1 16niPlayer + 4Ni−1Player +16Ni−1E(Player) otherwise,
(2) where the term 4Ni−1Player represents the number of frozen model parameters already trained in the 1 to i− 1 stages, the term 16niPlayer indicates the number of newly added model parameters with optimization states added in Procedure 2 and the term 16Ni−1E(Player) represents the number of PET parameters with optimization states added in Procedure 3.
Let L be the number of layers for the model that is finally obtained. Then, the solution of the following minimization problem can minimize the maximum memory requirement during the pre-training:
minimize {n1,...,nK}
{ max
i=1,...,K P STEPi
} s.t. L = NK (3)
Details of the discussion with specific examples are presented in Appendix B.

heading: 4 Experiments
text: This section demonstrates the effectiveness of the proposed method, STEP, through the pre-training experiments of LLMs. We investigate whether STEP can achieve a comparable validation perplexity to vanilla pre-training at the same computational cost. We also compared with ReLoRA (Lialin et al., 2024) as a conventional method of the parameterefficient pre-training method in a fair condition.

heading: 227M →352M 1024 18 →28

heading: 409M →668M 1760 11 →18

heading: 4.1 Datasets and Model
text: We used C4 (Raffel et al., 2020) as the training data and 10M tokens exclusively extracted from C4 as the validation data. We used the identical training data for all experiments.
The model configuration follows an architecture based on LLaMA (Touvron et al., 2023). The detailed configurations are shown in Appendix C. To confirm the differences in behavior due to model size, we selected three model sizes, namely, 352M, 668M, and 1.2B.

heading: 4.2 Configuration of STEP
text: We evaluated the effectiveness of STEP when the Growth Layer operator is applied once during its pre-training. This means that we set K = 2 in Equation 3 for STEP. Given the number of layers L with the fixed dimension of hidden layers, we compute n1 and n2 that can minimize the maximum memory requirements by Equation 3. Table 1 shows the calculated numbers of layers and parameters when the target model sizes are one of {352M, 668M, 1.2B}. Figure 2 shows an example of memory requirements when the target model size is 1.2B for vanilla pre-training and each stage of the two-stage STEP.
Layers are added to the upper part of the Transformer layers. The discussion about the position where layers are added is provided in Appendix E.

heading: 4.3 Results
text: Table 2 shows the validation es of vanilla pretraining, ReLoRA, and STEP. As shown in Table 2, STEP outperformed ReLoRA and achieved comparable validation to the vanilla pre-training while significantly reducing the maximum memory requirement from 5.6G to 3.6G (35.7% reduction), 10.7G to 6.6G (38.3% reduction), and 19.3G to 11.5G (40.4% reduction) for 352M, 668M, and 1.2B models, respectively. We also observed a desirable characteristic of increasing the model sizes, which led to a further reduction in maximum memory requirements, such as 35.7% to 40.4% for the 352M and 1.2B models, respectively. Based on these results, STEP has the potential to efficiently pra-train LLMs with reduced memory usage.

heading: 4.4 Ablation study
text: Type of Growth Layer Operators: We conducted an ablation study on Growth Layer Operators (Procedure 2) in STEP. We compared three Growth Layer operators: Stacking, Queueing, Interpolation-Copy, and Interpolation-Mean.
Stacking is proposed in Gong et al. (2019), which stacks additional layers. Queueing inserts new additional layers at the bottom. While the structure of layers resulting from Queueing is identical to that of Stacking, we need to consider this in STEP because PET is applied to the existing layers before Queueing. As in Gong et al. (2019), for both Stacking and Queueing, the weights of the additional layers are copied from the original layers. Interpolation-Copy and Interpolation-Mean
are Interpolation operators that use copy and mean initialization in Section 3.2, respectively. To simplify the discussion regarding the location of layer addition, the number of layers to be added is set to be the same as the total number of the model before Growing layers; that is, the number of layers in the model is doubled compared to before the addition.
The results of this ablation study are shown in Table 3. The performance in all settings is Interpolation-Mean > Interpolation-Copy > Queuing ≈ Stacking. One possible reason Interpolation outperformed Stacking and Queueing is that it can add layers to preserve the overall mechanism better. Several existing studies (Meng et al., 2022; Chen et al., 2024, 2023) have reported analysis results indicating that Transformers have distinct roles for the lower, middle, and upper layers, and it is thought that Interpolation can maintain this structure, resulting in better performance compared to other operators. Moreover, Interpolation-Mean outperformed Interpolation-Copy. This result suggests that the mean initialization is superior to copy initialization. More detailed experimental explanations are described in the Appendix D.

heading: 5 Conclusion and Limitation
text: Pre-training LLM requires substantial memory, posing a challenge for research. We proposed a new training method called STEP, which enables LLM pre-training with limited memory requirements. Our experiments demonstrated that STEP achieved comparative performance to vanilla LLM pre-training while minimizing peak memory usage.
Several limitations of our study should be addressed in future research. First, while we conducted experiments with up to two stages in STEP, the effectiveness of using more than three stages remains unexplored. Second, although our methods reduced memory usage, we did not observe significant enhancements in training speed. Third, although validation is the standard metric for evaluating the performance of pre-training, it is still unknown whether the models pre-trained by the proposed method can improve the downstream tasks. To investigate the downstream task, we need to fine-tune all the pre-trained models. Finally, our experiments focused on relatively smaller model sizes compared to the recent LLMs with billions of parameters, such as those with 7B or more.

heading: A About Zero Initialization
text: Shen et al. (2022); Wu et al. (2024) apply zeroinitialization to some modules when applying Interpolation to preserve the value. However, as Yao et al. (2024) points out, the existing layers may receive gradients similar to the previous stage, leading to unnecessary constraints and potentially slowing down the model’s convergence. Therefore, in this paper, we consistently refrained from using zero-initialization.

heading: B STEP with LLaMA and LoRA
text: In STEP, we use ReLoRA (LoRA) for PET and LLaMA as the model. When not considering Grouped Query Attention (Ainslie et al., 2023) in LLaMA, the Self-Attention layer contains four matrices of size (dhidden, dhidden). Additionally, the FFN layer has three matrices of size (83dhidden, dhidden), and there are two vectors of size dhidden for Layer Normalization. Therefore, Player is given by:
Player = 4d 2 hidden + 3(dhidden ×
8 3 dhidden) + 2dhidden
= 12d2hidden + 2dhidden (4)
Furthermore, since ReLoRA assigns two matrices of size (d, r) to a matrix of size (d, d), we have:
E(Player) = 8(rdhidden) + 3r(dhidden + 8
3 dhidden)
= 19rdhidden (5)
For example, if dhidden = 2048 and r = 128, equation 2 becomes, in units of GB,

heading: C Details of training configurations
text: The training configurations used in the experiment are shown in Table 4. These settings are the same across all experiments.
D Interpolation-Copy and Interpolation-Mean
To verify whether the mean initialization in Interpolation actually possesses the desired properties, we compare the cosine similarity between the output of the added layer ϕnew_i after Interpolation and the input to the layer ϕi+1 before Interpolation. In this case, if the output of ϕnew_i is similar to the input of ϕi+1, it can be considered that ϕnew_i appropriately processes the output from ϕi and outputs something that is easy for ϕi+1 to process. We apply Interpolation to a 334M model with nine layers that have been trained on the C4 dataset (Raffel et al., 2020) and has a perplexity of 18.54 on the validation dataset. Interpolation expands the model to a 668M model with 18 layers. Subsequently, using the same validation dataset, we obtain the embeddings of the output from ϕnew_i in the expanded 668M model and the embeddings of the input to ϕi+1 in the original 334M model before Interpolation. We then compare the cosine similarity between these two embeddings. The results are shown in Figure 3. As expected, using the average initialization yields a higher cosine similarity compared to copy initialization, suggesting that the
connections between layers are better preserved.

heading: E Effective position for adding new layers
text: This ablation study investigates the most effective position to add new layers when applying the Growth Layer operator using Interpolation-Mean in Procedure 2. In this ablation study, we perform layer additions on 409M →668M configurations in Table 1. We added the seven layers together starting from the bottom_index to the upper layer. In other words, when bottom_index is 3, the layers are added together at the top, and when it is 0, they are added at the bottom. The experimental results are shown in Table 5. As a general trend, the results indicate that adding layers to the upper part of the model leads to better performance improvements.

heading: Future Work
text: Several limitations of our study should be addressed in future research. First, while we conducted experiments with up to two stages in STEP, the effectiveness of using more than three stages remains unexplored. Second, although our methods reduced memory usage, we did not observe significant enhancements in training speed. Third, although validation is the standard metric for evaluating the performance of pre-training, it is still unknown whether the models pre-trained by the proposed method can improve the downstream tasks. To investigate the downstream task, we need to fine-tune all the pre-trained models. Finally, our experiments focused on relatively smaller model sizes compared to the recent LLMs with billions of parameters, such as those with 7B or more