Table of Contents
Fetching ...

Everybody Prune Now: Structured Pruning of LLMs with only Forward Passes

Lucio Dery, Steven Kolawole, Jean-François Kagy, Virginia Smith, Graham Neubig, Ameet Talwalkar

TL;DR

Bonsai introduces a forward-pass-only, gradient-free structured pruning framework for large language models, enabling memory-efficient pruning by estimating global module relevance via perturbation-based sub-model evaluations and a regression model guided by informative priors. The method supports iterative pruning across the entire model and achieves strong performance even under strict memory constraints, outperforming forward-only baselines and rivaling gradient-based approaches in several settings. Post-pruning adaptation (PPA) further enhances performance through distillation and fine-tuning, enabling competitive zero-shot capabilities even after substantial pruning. The work demonstrates practical implications for democratizing LLM compression on diverse hardware, while acknowledging runtime as a key tradeoff and outlining avenues for future improvements such as adaptive sampling and hybrid optimization strategies.

Abstract

Structured pruning is a promising approach to create smaller, faster LLMs. However, existing methods typically rely on backward passes, which can inflate memory requirements and compute costs. In this work we introduce Bonsai, a gradient-free structured pruning method that eliminates the need for backpropagation, significantly reducing memory requirements and compute costs while achieving state-of-the-art pruning performance. Bonsai uses forward-pass-only perturbative pruning to enable efficient compression of large models on a broader range of hardware configurations. Unlike existing structured pruning approaches, Bonsai not only achieves better compression with fewer resources, but also produces models that are twice as fast as those generated by semi-structured pruning. As a concrete demonstration, we use Bonsai to prune an 8B LLaMA-3 model to 50% sparsity on a single A6000 GPU -- a task infeasible with backprop-based methods, which require 2-3x memory. Our results show that removing backprop as a requirement not only enables pruning larger models on constrained hardware but can also lead to state-of-the-art efficiency and performance.

Everybody Prune Now: Structured Pruning of LLMs with only Forward Passes

TL;DR

Bonsai introduces a forward-pass-only, gradient-free structured pruning framework for large language models, enabling memory-efficient pruning by estimating global module relevance via perturbation-based sub-model evaluations and a regression model guided by informative priors. The method supports iterative pruning across the entire model and achieves strong performance even under strict memory constraints, outperforming forward-only baselines and rivaling gradient-based approaches in several settings. Post-pruning adaptation (PPA) further enhances performance through distillation and fine-tuning, enabling competitive zero-shot capabilities even after substantial pruning. The work demonstrates practical implications for democratizing LLM compression on diverse hardware, while acknowledging runtime as a key tradeoff and outlining avenues for future improvements such as adaptive sampling and hybrid optimization strategies.

Abstract

Structured pruning is a promising approach to create smaller, faster LLMs. However, existing methods typically rely on backward passes, which can inflate memory requirements and compute costs. In this work we introduce Bonsai, a gradient-free structured pruning method that eliminates the need for backpropagation, significantly reducing memory requirements and compute costs while achieving state-of-the-art pruning performance. Bonsai uses forward-pass-only perturbative pruning to enable efficient compression of large models on a broader range of hardware configurations. Unlike existing structured pruning approaches, Bonsai not only achieves better compression with fewer resources, but also produces models that are twice as fast as those generated by semi-structured pruning. As a concrete demonstration, we use Bonsai to prune an 8B LLaMA-3 model to 50% sparsity on a single A6000 GPU -- a task infeasible with backprop-based methods, which require 2-3x memory. Our results show that removing backprop as a requirement not only enables pruning larger models on constrained hardware but can also lead to state-of-the-art efficiency and performance.
Paper Structure (37 sections, 9 equations, 6 figures, 16 tables, 1 algorithm)

This paper contains 37 sections, 9 equations, 6 figures, 16 tables, 1 algorithm.

Figures (6)

  • Figure 1: Perplexity versus inference speed-up of pruned models for methods that use only forward passes on the parent model. At any given target perplexity, Bonsai produces the fastest model, resulting in improved latency and throughput. Circle sizes $\propto$ model's memory footprint.
  • Figure 2: Overview of Bonsai. Perturbations to the original model are generated according to an intelligent prior. These perturbed models are evaluated on downstream task and the collected data is used to build a simple linear model of module importances. Bonsai iteratively drops modules of least importance until the target model size is reached.
  • Figure 3: Perplexity at 50% sparsity of LLaMA-{1,2}-7B on Wikitext-2 and C4 datasets.
  • Figure 4: LLaMA-2 7B @ 50% sparsity (No PPA). Pertubation and regression components are both needed to make Bonsai effective. Experiment details in Appendix \ref{['appendix:no_regress_perturb']}.
  • Figure 5: LLaMA-2 7B pruned to 50% sparsity. See Appendix \ref{['appendix:prior_impact']} for experiment details and definitions of Wanda and Activation Magnitude priors.
  • ...and 1 more figures