Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning

8. Continue pre-training a small pre-trained model

Thanks for your insightful question! Our study was constrained to a 50B token budget due to computational limitations. Nevertheless, as depicted in Figure 1 of our paper, the trend of continued pre-training reveals no sign of convergence in performance. This strongly suggests that increasing the data scale further would continue to enhance the model."

Regarding your second question, further finetuning Open-LLaMA-3B-v2 for an additional 1T tokens was beyond our experimental capacity. Instead, we focused on smaller-scale Pythia models, each trained with 300B tokens. We pruned the Pythia-410M model down to 160M and continued its pre-training with 33B tokens (Sheared-Pythia-160M). This was compared against continuing the pre-training of Pythia-160M with an equivalent token amount. The revised Figure 11 shows that while Pythia-160M (300B+33B) initially outperforms, Sheared-Pythia-160M (300B+33B) learns more rapidly, eventually surpassing Pythia-160M. These preliminary results hint at a potential for greater performance gains from Sheared-Pythia with increased computational resources.

We think that pruning a larger model would yield more substantial benefits. This could be supported by findings from a previous study [1], which demonstrated that larger models perform better when pruned than training small models from scratch because larger models are easier to optimize. However, we will defer a deeper analysis of this hypothesis on the LLM scale to future work!

[1] Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers

We are greatly encouraged by the reviewer's positive feedback regarding the practicality and reproducibility of our work. In response to the reviewer's comments, we provide the following responses:

1. Specific to Transformer architecture

Transformer architectures have become ubiquitous in both vision (ViT) and language (LLMs) models, delivering state-of-the-art results on a wide range of applications. Their widespread relevance underscores the need for specialized pruning techniques for Transformer-based models. Moreover, considering the high training costs of large-scale Transformer-based LLMs – more than any other neural network – developing efficient training and compression methods for these transformers is vital.

2. Justification for the generality of our approach

Thanks for your input! We would like to argue that our method has broader applicability for several reasons:

• The landscape of open-source LLMs is evolving at a blazingly fast pace. In the recent two months, many strong open-source models have been released, e.g., Mistral-7B (the strongest 7B model so far), Yi-34B (the strongest 34B model so far), DeepSeek-Coder (the strongest open-source code model so far with 1.3B, 6.7B, and 33B). However, such models are usually not released across many scales, largely due to the expensive training cost. Our method, instead, provides an efficient general approach to derive smaller-scale foundation models from any existing models, including all the ones that I mention, to any size and configuration. We believe our work conveys an important message: to pre-train a new model, we should structurally prune from existing models instead of training it from scratch (unless the target model is larger than all the existing models).
• The method is not limited to the strongest suite of models. In the general response, we show that pruning relatively weaker models (Pythia models) also accelerates pretraining a smaller model.
• The method could also be useful for developing domain-specific models as the reviewer pointed out.

In summary, our approach presents a compelling strategy for expediting the development of open-source large language models (LLMs) at various scales, and most importantly is affordable to communities with limited computational resources.

3. Need to take 2T pretraining data into account

Thank you for your input! Our approach is grounded in the realistic context where strong large language models (LLMs) are readily available, yet typically not in a diverse range of scales, particularly the latest releases, e.g., the strongest 7B model Mistral only has one single size. Our paper underscores the strategic use of these existing assets to derive smaller models at varied scales. While the pre-training expense of the source model is significant, it's essential to recognize that:

• this often represents a sunk cost, usually borne by well-resourced institutions (as exemplified by our non-requirement to train the publicly available LLaMA-7B).
• this cost can be effectively spread over the process of pruning for different scale models, thereby optimizing resource utilization.

4. Compute for calculating the reference loss

In terms of acquiring reference losses, we ran inference on 14M tokens with 3 models (Llama2-7b, 13b, and 70b). And the whole process took less than 4 A100 GPU hours. The cost to compute the reference losses is negligible compared to pruning (training on 0.4B tokens) and dynamic batch loading (training on 50B tokens). We have added this information to our paper.

5. Pruning another source LLM

Thanks for your suggestion! We conducted similar experiments on Pythia, and please refer to the general response, 1. Pruning a weaker and smaller foundation model, for more details.

6. LLaMA’s pretraining data and RedPajama

Thanks for capturing this – updated in our revision!

7. Placement of figure 8

Thanks for the comment and we updated the draft!

Dear Reviewer urdF,

We’d love to hear your thoughts on our newly added results and hopefully they have addressed your concerns for the paper! We really appreciate your efforts in providing insightful comments and suggestions to us!

We are encouraged that the reviewer acknowledged that the Sheared-LLaMA models underwent thorough evaluations and demonstrated strong performance. Thanks for your comments and suggestions on the work!

1. Concerns about the novelty of the paper.

We disagree with this point, as our proposed methods of targeted structured pruning and dynamic batch loading are both different from previous work significantly, which we clarify below:

Targeted structured pruning:

• The goal of targeted structured pruning is to achieve an exact target transformer architecture that is pre-specified based on a source transformer model. To the best of our knowledge, none of the works mentioned by the reviewer [1, 2, 3, 4, 5] have done this. This is based on the intuition that existing models’ configurations have already been well optimized to balance model expressivity and inference efficiency, and empirically we showed that it is more effective than pruning to non-uniform structures with only sparsity specified.
• Another important contribution is that we prune a general-purpose, large language model with billions of parameters (7 billion parameters in our experiments), a scale that was not demonstrated in the literature. The implications of pruning at such a large scale are not well understood, and our work serves as one to provide such insights.

Dynamicbatch loading:

• Although our dynamic loading method is inspired by Doremi [6], they differ substantially in that Doremi requires a complicated multi-staged process to train a reference model and a proxy model before the final training run. Our dynamic batch-loading algorithm does not require training any additional models, making it significantly more efficient and easier to use than Doremi.
• The other contribution is that we introduce a reference loss derived from an estimate of scaling law, which takes scale differences into account, and is a more principled solution to estimating a reference loss.

We have updated the paper to better highlight our contributions and differences with prior works in Section 2.1 and 2.2.

2. Fair comparison to LLM-Pruner

Please refer to the general response, 1.A fair comparison compared to LLM-Pruner, where we provided a fair and thorough comparison to LLM-Pruner in terms of loss, model architecture, inference speed and training throughput.

3. Pruning a weaker and smaller foundation model

Thanks for your suggestion! We conducted additional experiments on Pythia-440M models, and please refer to the general response, 2. Pruning a weaker foundation model, for more details.

Dear Reviewer wgpN,

As it’s the last day for discussion, we wonder if our responses have addressed your concerns for the paper. We would really love to have your inputs on it and are happy to provide more information if needed!

Thank you for the valuable feedback! We appreciate that you acknowledge the efficiency implications of our proposed approach to pre-training and the strong performance of our models. We address your comments below:

1. A far comparison to LLM-Pruner

We compared our pruning approach to LLM-Pruner, both without dynamic batch loading. Please refer to the general response, ** A fair comparison compared to LLM-Pruner**, where we provided a fair and thorough comparison to LLM-Pruner in terms of loss, model architecture, inference speed and training throughput. Besides, the PPL difference of 0.24 roughly corresponds to 1 point difference in downstream performance on the training trajectory, and we do not consider it as marginal.

2. Dynamic batch loading

The core contribution of dynamic batch loading is exactly to automatically figure out what the “balanced data” should be for both pruning and continued pre-training, without adding additional compute overhead.To add additional contexts: In our initial experiments, we also explored directly excluding “easy domains” like GitHub, StackExchange, and ArXiv during the pruning stage. However, it does not address the loss imbalance issue. Please refer to our response to reviewer qt4Y if you are interested in more details. We also explored pruning with more tokens but generally observed that the loss imbalance issue persists. Hence the loss imbalance problem is not likely to be solved by putting in more pruning budget. [1, 2] are about different pruning techniques. We experimented with [1], and still observe a similar issue with loss disparity across domains.

In conclusion, we argue that such a loss imbalance problem cannot be easily solved by manually curating data, adding more pruning budgets, or switching pruning methods. We believe our dynamic batch loading is essential, and its strength is demonstrated in our ablations in Section 4.1.

3. Comparison to CoFiPruning

We recognize that Sheared-LLaMA does not outperform CoFiPruning in perplexity when both methods use the same amount of data for pruning. However, a key difference is that CoFiPruning produces a non-uniform model, which introduces extra inference costs.

In practice, the perplexity gap could be compensated by slightly training Sheared-LLaMA longer for about 0.6B tokens in our experiment. Since this extra training is a one-time cost and inference speed is a crucial factor for LLM serving, the choice to prune towards a uniformly layered architecture is justified for long-term efficiency.

4. Typo

Thank you for catching the typo; we've updated the paper. The LLM-Pruner has 1.6B parameters, 0.3B more than Sheared-LLaMA due to its lack of support for pruning hidden dimensions, including word embeddings. While embeddings have minimal inference overhead, we've maintained parity in the remaining model parameters between Sheared-LLaMA and LLM-Pruner for a fair comparison.

Dear Reviewer bT2X,

We hope our newly added experiments help address your concerns for the paper! We would love to hear your thoughts about them and provide more information if needed.

Really appreciate it!

4. A version of Figure 4 with perplexity or loss.

In our revised draft, we added Figure 12 to illustrate the loss trends during continued pre-training of both the INCITE model and our pruned model. Notably, the INCITE model's loss remains relatively unchanged throughout this phase, whereas the loss of our pruned model progressively decreases, ultimately falling below that of the INCITE model. This trend is consistent with the downstream task performance depicted in Figure 4.

5. Typos

Thanks for spotting them our typos and we have fixed all of them in the revision.

6. Initialization of Lagrange multipliers

We randomly initialize the Lagrange multipliers with value 0. Additionally, we've incorporated the related work you mentioned – It’s nice to get to connect to classical works!

7. Explanation for non uniformity

Thanks for pointing this out! The GPU architecture and model parallelism used in training large models mean that non-uniform structures can increase memory allocation overhead and disrupt load balancing across GPUs. Additionally, varying shapes, even with the same number of elements, can significantly impact inference latency for matrix multiplication operations. We empirically demonstrate in Section 4.2 that uniform pruned structures outperform non-uniform ones in speed. We added these intuitive explanations to our revised draft.

8. adding perplexity after continued pretraining

For a fair comparison, we only included the pruned models’ perplexities without continued pre-training in Table 4. In the future, we will consider adding an experiment on continuing pre-training CoFiPruning to complete this comparison. Thanks for the suggestion!

9. Adding Bartoldson et al., 2023

Thanks for the nice pointer! We have incorporated it into our revision.

1. Adding inference speedup for different pruning approaches

Thanks for your suggestion. We've included below the inference speed (tokens/second) comparisons for all methods mentioned you mentioned. These experiments were all conducted on a single A100 80GB GPU, generating up to 2048 tokens with a batch size of 1.

Sheared-LLaMA is way more efficient at inference than LLM-Pruner, largely due to the better chose final architecture (See general response, ** A fair comparison compared to LLM-Pruner**, for a more thorough comparison). Wanda type of semi-structured pruning also achieves inference speedup compared to the source model. But it is not as fast as small dense models, and is restricted to 50% sparsity. Measuring throughputs (using the maximum batch size) will give Sheared-LLaMA even more advantage than this setup. We have added the analysis to Appendix H.

2. Pruning with CC, C4, Wiki and Books

Regarding point 1, in the project's initial phase, we experimented by pruning without including data from GitHub, StackExchange, and ArXiv (domains retaining more performance). We specifically pruned LLaMA1-13B to 7B using a mix of C4 (40%), CC (40%), Wiki (10%), and Books (40%) datasets. Post-pruning, we further pre-trained this model using the RedPajama dataset with all of the domains.

As shown in the table below, post-pruning perplexity differences were fairly uniform across domains. However, after continued pre-training with data from all the domains of RedPajama, the loss disparity grew: the loss difference on GitHub/StackExchange got much smaller than those on domains like C4. This indicates that simply excluding easier-to-recover domains during pruning doesn't address the loss disparity issue in continued pre-training.

These findings led us to develop dynamic batch loading, a more effective solution to tackle domain loss disparity. For a comprehensive understanding, we've detailed these initial experiments in Appendix G of our revised submission.

3. Use a larger set of validation data for more complex domains

Enlarging the validation set of complex domains for training stability is an excellent idea (point 2), which will be likely to provide improved calibration and stability for dynamic data loading. We did not explore this direction thoroughly in our initial experiments, so there is certainly room to build on this suggestion in future work to better understand the impacts!