PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, Danqi Chen

arXiv:2310.06694v2cs.CLcs.AIcs.LG

TL;DR

Training smaller, competitive LLMs from scratch remains computationally costly, motivating the use of existing larger models. The paper combines targeted structured pruning with dynamic batch loading and continued pre-training, producing smaller Sheared-LLaMA models that outperform similarly sized open-source models using only a fraction of standard pre-training compute.

  • Problem

    Training even small billion-parameter LLMs from scratch requires substantial computational resources, motivating smaller competitive models built from existing ones.

  • Method

    The approach combines targeted structured pruning to reach a specified architecture with dynamic batch loading and continued pre-training to improve data efficiency.

  • Results

    Sheared-LLaMA-1.3B and Sheared-LLaMA-2.7B outperform similarly scaled open-source LLMs on 11 downstream tasks and instruction tuning after using 50 billion additional tokens.

  • Takeaways & Limitations

    Structured pruning offers a promising path to producing competitive smaller-scale LLMs at low cost when strong large-scale models are available.

  • Takeaways & Limitations

    The method depends heavily on open-source pre-training datasets and models, and the experiments used only a 7B-parameter source model.

Abstract

from arXiv · show

The popularity of LLaMA (Touvron et al., 2023a;b) and other recently emerged moderate-sized large language models (LLMs) highlights the potential of building smaller yet powerful LLMs. Regardless, the cost of training such models from scratch on trillions of tokens remains high. In this work, we study structured pruning as an effective means to develop smaller LLMs from pre-trained, larger models. Our approach employs two key techniques: (1) targeted structured pruning, which prunes a larger model to a specified target shape by removing layers, heads, and intermediate and hidden dimensions in an end-to-end manner, and (2) dynamic batch loading, which dynamically updates the composition of sampled data in each training batch based on varying losses across different domains. We demonstrate the efficacy of our approach by presenting the Sheared-LLaMA series, pruning the LLaMA2-7B model down to 1.3B and 2.7B parameters. Sheared-LLaMA models outperform state-of-the-art open-source models of equivalent sizes, such as Pythia, INCITE, OpenLLaMA and the concurrent TinyLlama models, on a wide range of downstream and instruction tuning evaluations, while requiring only 3% of compute compared to training such models from scratch. This work provides compelling evidence that leveraging existing LLMs with structured pruning is a far more cost-effective approach for building competitive small-scale LLMs

1 INTRODUCTION

The paper asks whether existing pretrained LLMs can be converted into smaller competitive models with much less compute than training from scratch. It proposes targeted structured pruning and dynamic batch loading, then demonstrates competitive Sheared-LLaMA models at 1.3B and 2.7B parameters.

  • Training even small billion-parameter LLMs from scratch requires substantial computational resources, motivating smaller models built from existing pretrained systems.
  • Structured pruning is used to develop smaller general-purpose LLMs, despite typical performance degradation when pruning lacks substantial post-pruning compute.
  • Targeted structured pruning searches a source model’s substructures to produce a specified target architecture while maximizing preserved performance under constraints.
  • Dynamic batch loading samples training data from each domain according to its rate of loss reduction, improving data efficiency and accelerating performance improvement.
  • Figure 1 reports that Sheared-LLaMA-2.7B surpasses similarly scaled open-source models while using 1/32 (3%) of the budget to achieve on-par performance with OpenLLaMA-3B-v2.
  • Sheared-LLaMA-1.3B and Sheared-LLaMA-2.7B, pruned from LLaMA2-7B, outperform similarly sized Pythia, INCITE, and OpenLLaMA models across 11 downstream tasks and instruction tuning.

2 LLM-SHEARING

LLM-shearing is framed as a two-stage process: prune a source model into a specified compact architecture, then continue pre-training it to recover performance. The method combines target-shape structured pruning with dynamic batch loading that adjusts domain proportions according to loss behavior.

  • Two-stage LLM-shearing: Continued pre-training is crucial after pruning because pruning reduces parameters but inevitably causes a performance drop.The pruning stage is followed by language-model pre-training of the finalized compact model.
  • Targeted structured pruning: Target architectures are selected from existing pre-trained model configurations to balance model expressiveness and inference efficiency.The method aims to scale down all transformer dimensions to a specified target configuration.
  • Targeted structured pruning: Learned masks control pruning at layer, hidden-dimension, attention-head, and intermediate-dimension granularities.The masks support both global structures, such as layers and hidden dimensions, and local structures, such as heads and intermediate dimensions.
  • Targeted structured pruning: The pruning objective jointly optimizes model weights and masks while enforcing a pre-specified target architecture.Hard-concrete mask distributions enable discrete prune-or-retain decisions, and the language-modeling loss is computed with masked weights.
  • Dynamic batch loading: The method compares domain losses with reference losses predicted by a scaling function or taken from the source model.The scaling reference performs slightly better downstream, especially on math and coding tasks, while the source reference remains viable when source-model scaling data is unavailable.
  • Dynamic batch loading: Dynamic batch loading reallocates training data across domains according to their loss reduction so domains approach reference losses at roughly similar times.It is applied during both pruning and continued pre-training, using validation losses and dynamically updated data proportions.

3 EXPERIMENTS

The experiments evaluate Sheared-LLaMA models against similarly sized language models on downstream tasks and instruction tuning, using RedPajama-based pruning and continued pre-training. Sheared-LLaMA outperforms comparable open models despite a limited training budget.

  • 3.1 SETUP: The study prunes LLaMA2-7B into 2.7B and 1.3B parameter target models and compares them with several similarly sized baselines.Baselines include OPT, Pythia, TinyLlama, INCITE-Base, and OpenLLaMA models.
  • 3.1 SETUP: RedPajama supplies seven domains for pruning and continued pre-training, with 0.4B tokens allocated to pruning and 50B tokens to continued pre-training.The domains include CommonCrawl, C4, Github, Wikipedia, Books, ArXiv, and StackExchange.
  • 3.2 SHEARED-LLAMA OUTPERFORMS LMS OF EQUIVALENT SIZES: Sheared-LLaMA models outperform existing similarly sized models on standard downstream benchmarks despite using approximately 50B tokens for pruning and continued pre-training.The comparison includes models trained with substantially larger compute budgets.
  • 3.2 SHEARED-LLAMA OUTPERFORMS LMS OF EQUIVALENT SIZES: Sheared-LLaMA-1.3B outperforms TinyLlama-1.1B, although TinyLlama-1.1B was pre-trained on 3T tokens.The paper presents this as evidence of greater sample efficiency for training smaller-scale LLMs.
  • 3.2 SHEARED-LLAMA OUTPERFORMS LMS OF EQUIVALENT SIZES: Instruction-tuned Sheared-LLaMA achieves higher win rates than comparable pre-trained models, including Pythia-1.4B, INCITE-Base-3B, and OpenLLaMA variants.Instruction tuning uses instruction-response pairs sampled from ShareGPT.

4 ANALYSIS

The analysis examines dynamic batch loading, structured-pruning architecture choices, and compute allocation. Dynamic loading balances domain-wise loss reduction and improves downstream performance, while targeted structures improve inference efficiency and more pruning-stage compute improves perplexity.

  • 4.1 EFFECTIVENESS OF DYNAMIC BATCH LOADING: Dynamic batch loading is evaluated through final loss, domain data usage, and downstream task performance for Sheared-LLaMA-1.3B.The analysis compares dynamic loading with the original RedPajama distribution.
  • 4.1 EFFECTIVENESS OF DYNAMIC BATCH LOADING: Dynamic batch loading reduces variation in domain loss differences relative to the estimated reference loss.Original loading leaves GitHub below the reference loss while C4 lags behind.
  • 4.1 EFFECTIVENESS OF DYNAMIC BATCH LOADING: Dynamic batch loading uses more Book and C4 data, indicating that these domains are more challenging for a pruned model to recover.Table 3 compares dynamic domain usage with the original RedPajama proportions.
  • 4.1 EFFECTIVENESS OF DYNAMIC BATCH LOADING: Pruned models trained with dynamic batch loading achieve better downstream performance than models trained with the original RedPajama distribution.The analysis links this improvement to more balanced loss reduction.
  • 4.2 TARGETED STRUCTURED PRUNING: Targeted structured pruning produces uniform layer configurations, which provide faster inference than non-uniform CoFiPruning architectures.The comparison reports validation perplexity and inference speed in tokens per second on an Nvidia A100 GPU.
  • 4.2 TARGETED STRUCTURED PRUNING: Given the same compute budget, sparsity level, and training data distribution, targeted structured pruning achieves lower perplexity, a more optimized architecture, and faster inference than LLM-Pruner.The comparison is reported in Appendix F.2.
  • 4.3 ADDITIONAL ANALYSIS: Within a fixed 5B-token budget, increasing the pruning allocation consistently improves perplexity, but the authors allocate 0.4B tokens to pruning because pruning is more expensive than continued pre-training.The trade-off is evaluated across pruning and continued pre-training budgets.

5 RELATED WORK

Prior work studies pruning, efficient pre-training, and data selection, but general-purpose LLMs are less redundant than task-specific models and pruning often causes substantial performance drops. The paper positions pruning with continued training as an efficiency-oriented alternative to several existing approaches.

  • Pruning: Structured pruning has commonly compressed overparameterized task-specific models with limited performance impact.The related work covers pruning methods in computer vision and natural language processing.
  • Pruning: For general-purpose LLMs, unstructured, semi-structured, and structured pruning can cause significant performance drops even at modest sparsity.Previous approaches often fix the original models or tune them only minimally.
  • Efficient pre-training approaches: Efficient pre-training research includes quantization, knowledge distillation, dynamic architectures, and efficient optimizers.These approaches target training or inference efficiency through different mechanisms.
  • Efficient pre-training approaches: The paper describes pruning and continued training as more cost-effective than teacher-based distillation for pre-training LLMs because distillation incurs additional teacher inference cost.This comparison is presented as orthogonal to the pruning approach.
  • Data-based approaches: Data-based efficiency methods include deduplication, loss-based batch selection, reducible-loss prioritization, and proxy-model data-mixture optimization.These methods adjust which data are retained, selected, or weighted during training.

6 DISCUSSION

The discussion presents structured pruning with dynamic batch loading as a cost-efficient route to competitive smaller LLMs, while noting dependence on available data and source models. The pruning-mask framework learns sparse structures and uses scaling laws to estimate reference losses for target models.

  • Structured pruning depends heavily on open-source pre-training datasets and models, and uncovered domains may not recover performance well.
  • The method was evaluated only with a 7B-parameter source model because of computational constraints, although the authors describe it as generalizable to larger models.
  • Hard concrete masks continuously relax binary pruning decisions, enabling the model to learn sparse structures during training.
  • Masks are trained with Lagrange multipliers and then thresholded so removing masked parameters produces the target model shape.
  • The scaling law models language-model loss as a function of model size N and dataset size D, with scaling factors estimated separately for each domain.
  • Using LLaMA2-7B, 13B, and 70B checkpoints yields predicted losses for hypothetical smaller models, but limited data points may bias the LLaMA2-2.7B estimate.

C TRAINING DETAILS

The training details specify parallelized pre-training and instruction-tuning procedures for Sheared-LLaMA and baseline models. Evaluation uses ShareGPT instructions, GPT-4 pairwise judging, and example generations from the resulting models.

  • Training setup: Training uses fully sharded data parallelism, FlashAttention V1, and a cosine learning-rate schedule decaying to 10% of the peak value.
  • Model configurations: Sheared-LLaMA configurations closely follow Pythia-1.4B and INCITE-Base-3B, with intermediate-size differences resulting from its LLaMA2-7B lineage.
  • Instruction tuning: Instruction tuning uses 10,000 ShareGPT instruction-response pairs, while evaluation samples another 1,000 instructions.
  • Instruction tuning: GPT-4 evaluates response pairs twice with swapped presentation order, and the reported score averages both rounds to reduce position bias.
  • Instruction tuning: Each instruction is prepended with a fixed helpful-assistant prompt before instruction-tuning evaluation.
  • Instruction-tuning results: After instruction tuning, Sheared-LLaMA-2.7B produces long, coherent, and informative responses in sampled examples.

F.1 DATA USAGE IN CONTINUED PRE-TRAINING

Dynamic batch loading changes domain sampling during continued pre-training according to evolving losses. Domain weights stabilize after about 30% of training, while Wikipedia shows an early spike and other domains change steadily.

  • Domain weights stabilize after approximately 30% of training under dynamic batch loading.
  • Wikipedia exhibits an anomalous early spike in loading weight, unlike the steady monotonic changes observed for the remaining domains.
  • The analysis tracks both each batch’s domain weights and cumulative data usage across domains.
  • Table 10 provides example instruction-tuned outputs from Sheared-LLaMA-1.3B and Sheared-LLaMA-2.7B.

F.2 COMPARISON TO LLM-PRUNER

Against LLM-Pruner and continued-pretraining baselines, targeted structured pruning yields better loss, more practical architectures, and faster training or inference in several comparisons. Results also show advantages for LLaMA2 pruning, while math and coding remain weaker areas and some findings are preliminary.

  • Comparison to LLM-Pruner: At roughly 1.23B parameters excluding embeddings, targeted structured pruning is compared fairly with LLM-Pruner using equal pruning and continued-pretraining token budgets.
  • Comparison to LLM-Pruner: With the same amount of consumed data, targeted structured pruning achieves lower loss than LLM-Pruner.
  • Comparison to LLM-Pruner: LLM-Pruner produces an unconventional architecture with intermediate size smaller than hidden size because it cannot prune hidden dimensions or layers.
  • Comparison to LLM-Pruner: Sheared-LLaMA architectures provide higher throughput for both inference and continued training than LLM-Pruner models, whose training throughput is around 60% of Sheared-LLaMA’s.
  • Comparison to LLM-Pruner: At the same parameter scale, the proposed method combines lower perplexity, a more reasonable architecture, and faster inference speed.
  • Dynamic-loading references: Scaling-reference dynamic loading consistently achieves slightly better downstream performance than source-reference loading.
  • Pythia pruning: Sheared-Pythia-160M initially trails continued-pretraining Pythia-160M but learns faster and eventually exceeds it; the authors characterize these results as preliminary.
  • LLaMA1 versus LLaMA2: Pruning from LLaMA2 consistently outperforms pruning from LLaMA1, although the performance difference is not very large.

F.7 COMPARISON TO FURTHER CONTINUAL PRE-TRAINING INCITE-BASE-3B

Sheared-LLaMA-2.7B improves rapidly during continued pre-training and surpasses INCITE-Base-3B, although the advantage of pruning depends on the relative strength of the available small model. Domain-specific loss disparities also motivate dynamic batch loading.

  • Sheared-LLaMA rapidly improves during continued pre-training and surpasses INCITE-Base-3B, which starts with higher task accuracy but later plateaus.The comparison uses loss curves and downstream performance between similarly sized models.
  • Pruned models from a strong base model can provide a better initialization for continued pre-training than an existing LLM of equivalent size.This conclusion is drawn from the comparison between Sheared-LLaMA-2.7B and continually pre-trained INCITE-Base-3B.
  • Existing small models may be better starting points when they are competitive with the pruning source model.The advantage of pruning is greater when the performance gap between the source and small model is larger.
  • Dynamic batch loading: Excluding domains with rapidly decreasing losses makes perplexity differences more even after pruning, but disparities grow again during continued pre-training on all domains.The experiment used GitHub, StackExchange, and ArXiv exclusions during pruning and evaluated seven-domain continued pre-training.
  • Dynamic batch loading: These domain-specific loss disparities motivated dynamic batch loading as a more principled way to adjust training data composition.The supplied passage identifies dynamic batch loading as the response to the observed imbalance.

F.9 INFERENCE SPEED ANALYSIS

The inference-speed analysis compares dense Sheared-LLaMA models with the LLaMA2-7B source, semi-structured Wanda pruning, and LLM-Pruner. Sheared-LLaMA is faster than LLM-Pruner, while Wanda's speedup is constrained by its sparsity requirement.

  • Sheared-LLaMA achieves higher inference speed than LLM-Pruner because its resulting architecture is more optimized.Both approaches are compared at the same number of nonembedding parameters.
  • The analysis measures generation speed in tokens/second on an A100 GPU with up to 2048 generated tokens and batch size 1.The compared models include LLaMA2-7B, Sheared-LLaMA-1.3B, Sheared-LLaMA-2.7B, Wanda variants, and LLM-Pruner.
  • LLM-Pruner produces an architecture whose intermediate size is smaller than its hidden size, contrary to the stated transformer design convention.The passage links this architectural difference to the comparison of inference speed.
  • Wanda-type semi-structured pruning speeds inference relative to the source model but remains slower than small dense models.Its speedup is feasible only at 50% sparsity, making it less flexible.

G FREQUENTLY ASKED QUESTIONS

The FAQs clarify how to interpret Sheared-LLaMA’s token budget, distinguish dynamic batch loading from Doremi, and discuss choosing a pruning source model. They also state that source-model selection is difficult and rests partly on a simplified information-counting thought experiment.

  • Token budget: TinyLlama-1.1B, trained on 3T tokens, has downstream benchmark performance comparable to Sheared-LLaMA-1.3B when both are used as base models.TinyLlama provides a later reference point for models trained at a similar scale.
  • Dynamic batch loading: Dynamic batch loading and Doremi both adjust domain data distributions using model loss and exponential ascent, but dynamic batch loading requires no reference or proxy model.Dynamic batch loading derives a reference loss from a fixed evaluation set and can be applied directly to any model.
  • Source-model size: Choosing the optimal source-model size for pruning is challenging.The paper uses a simplified uniform-information thought experiment to compare source models trained with the same token count.
  • Source-model size: Under the thought experiment, pruning a 7B model trained on 2T tokens to 1.3B retains an estimated 0.37T tokens of information, versus 0.2T from a 13B model trained on the same token count.The estimate assumes information is uniformly distributed across parameters, an assumption the paper notes may not hold.
Loading 2310.06694v2…