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

Abstract: Pre-training large language models (LLMs) faces significant memory challenges due to the large size of model weights. We introduce STaged parameter-Efficient Pre-training (STEP), which integrates parameter-efficient tuning techniques with model growth. We conduct experiments on pre-training LLMs of various sizes and demonstrate that STEP achieves up to a 53.9% reduction in maximum memory requirements compared to vanilla pre-training while maintaining equivalent performance. Furthermore, we show that the model by STEP performs comparably to vanilla pre-trained models on downstream tasks after instruction tuning.

## introduction large language models (llms) have become an indispensable foundational technology in artificial intelligence. recent llm development trends, based on scaling laws @xcite , involve pre-training transformer models with a vast number of parameters on massive datasets @xcite . consequently, the pre-training of llms requires substantial computational resources, typically involving thousands of gpus @xcite . this enormous computational demand presents a significant obstacle to llm research. to tackle this challenge, we consider methods for reducing the computational demand in llm pre-training. while there are various approaches to reducing this, we introduce a pre-training method that maintains performance equivalent to vanilla pre-training while constraining the maximum gpu memory requirements to a predetermined threshold. specifically, our approach combines model growth @xcite through layer addition with parameter-efficient tuning techniques @xcite , which are commonly used in fine-tuning. for a detailed explanation of the proposed method, figure 1 presents an overview of our procedure. our approach formulates the maximum memory requirements for each stage of the sequential model growth as an integer programming problem, using model configurations as variables. we solve this optimization problem to determine the optimal model configurations for each stage, thereby controlling model growth settings to minimize peak memory usage prior to pre-training execution. this approach enables pre-training while maintaining memory requirements within a predetermined threshold. hereafter, we refer to our method as staged parameter efficient pre-training (step). we demonstrate that step achieves up to a 53.9% reduction in maximum memory requirements compared to vanilla pre-training while maintaining equivalent perplexity and performance on domain-specific tasks. furthermore, we verify that step does not negatively affect the performance of downstream tasks by demonstrating that steped models perform on par with the vanilla pre-trained model. ## related work several memory-efficient training approaches have been actively developed in the literature of training llms @xcite . one of the primary approaches involves reducing the number of trainable parameters. notable examples include parameter-efficient tuning (pet) methods such as adapter @xcite and @xcite . meanwhile, to reduce flops during pre-training, model growth techniques have been proposed @xcite , where training begins with a small-scale model and continues as the model parameters are gradually expanded. our proposed method aims to achieve memory-efficient pre-training by appropriately combining pet and model growth techniques. paremeter-efficient tuning. pet has primarily been developed for fine-tuning llms. for instance, lora is a technique that adds new adapters (lowrank matrices) while keeping the pre-trained llm parameters frozen, and only trains these adapters. since adapters typically contain few parameters, training can be accomplished with minimal memory requirements. pet is now being applied to pre-training applications. here, we describe two representative methods: relora @xcite and ga-lore @xcite . relora is a method for pre-training llms using lora. a distinctive feature of relora is that it begins with vanilla pretraining and transitions to lora during the training process. consequently, from a peak memory requirement perspective, relora requires the same amount of memory as vanilla pre-training. galore is a method that leverages the low-rank structure of gradients to reduce optimizer states while maintaining performance equivalent to vanilla pre-training. unlike relora, galore operates with low memory requirements throughout the entire training process. these methods can reduce memory usage compared to vanilla pre-training, but they slightly underperform. growing pre-trained model. recent studies have shown that growing a smaller model and then continuing to train the larger model can achieve comparable performance with fewer flops compared to training a large model from scratch @xcite @xcite . in these methods, the operation of increasing the model size is called the growth operator, expanding the dimensions of transformer @xcite layers and adding new layers. since existing studies train the full parameters of the model, this approach does not reduce the maximum memory requirements. ## ablation study we examine the effective position for new layers and the effectiveness of pet, both key components of step. 8 we used the model settings with a target size of 680m from as shown in the effect of pet parameters. this experiment verifies whether the pet introduced in step contributes to performance improvement. specifically, we conducted an experiment skipping procedure 3 in section 3.1. as shown in ## conclusion pre-training llm requires substantial memory, posing a challenge for llm research. we proposed a novel training method called step, which enables llm pre-training with reduced memory requirements. our experiments demonstrated the effectiveness of step; specifically, step achieved equivalent performance to vanilla pre-training and downstream tasks after instruction tuning, while reducing peak memory usage by up to 53.9%. we hope our results encourage researchers who aim to engage in llm pre-training research but have only limited computing resources.