4 Building smaller and faster LLMs with depth pruning

This chapter introduces depth pruning as a way to make large language models smaller and faster by removing entire Transformer blocks. It explains why this aggressive form of pruning is effective: beyond reducing parameters, it cuts memory traffic, lowers KV-cache growth, and can improve inference speed and memory use, especially in deep models where some redundancy exists across blocks.

The chapter then contrasts simple static selection rules with data-driven block selection. Static methods rely on position or weight magnitude and are easy to apply, but they ignore how a model actually behaves on a task. To overcome that, the chapter uses PyTorch hooks to inspect block inputs and outputs, then measures how much each block changes activations with cosine similarity. This yields a task-specific importance score, calibrated on WinoGrande, so blocks can be ranked by how much they matter for the data being processed.

Finally, the chapter evaluates different pruning choices on benchmarks and shows that the “best” blocks to remove depend on the task and calibration data. It finds that preserving early and late blocks is often useful, but that data-driven pruning can outperform purely static rules, especially when combined with simple protection heuristics. The chapter also connects these practical results to the ideas in “Shortened LLaMA” and “ShortGPT,” highlighting that depth pruning is not only scalable and efficient, but also supported by recent research as a strong alternative to more expensive pruning approaches.

Structure of the transformer blocks of the DistilGPT2 and Qwen3-06B models. Each block contains two functional modules: Attention and MLP. In turn, each module is composed of individual computational layers such as linear projections, normalizations, and activation functions.

Impact of block selection on depth pruning. The gain in inference (+15%) is identical as we are removing four blocks in both scenarios. The performance cost varies drastically depending on which ones are removed.

Comparison of block selection strategies for depth pruning. From left to right: The full base model. Static strategies that remove blocks based on their position: the first, the last, or the middle ones. A data-driven strategy, where the removed blocks are scattered according to their actual contribution to the task. The data-driven strategy combined with a protection heuristic, which preserves the initial and final blocks.

Three pairs of vectors from left to right: Vectors with similar meaning point in similar directions, unrelated vectors in different directions, vectors with opposite meanings in opposite directions.

In the block evaluation pipeline, hooks are registered, a forward pass of the model is executed using the dataset, and the activations are captured to calculate the block importance. The result is a list containing the block identifier with the assigned importance.

Layer importance across all 28 blocks of Qwen3-0.6B: static magnitude (weight-based) versus data-driven cosine similarity calibrated on Winogrande. Both metrics are normalized to the [0,1] range. Block 0 dominates the cosine ranking by a wide margin.

• Depth pruning eliminates entire Transformer blocks from a model. While aggressive, this approach targets the real bottleneck in LLM inference: data movement between memory and compute units, not just computational operations.
• The primary benefit of depth pruning is reducing memory-bound operations. Each removed block eliminates a complete round of loading data from VRAM, processing it, and writing results back, directly improving inference latency and throughput.
• Static block selection strategies (position-based or magnitude-based) are simple and computationally free but ignore how the model actually processes your specific data. The common heuristic of preserving the first four and last two blocks provides some safety but remains a general rule.
• PyTorch hooks allow you to intercept and capture the input and output activations of any module during inference without modifying the model's code. Use register_forward_hook() to attach a function that receives the module, input, and output tensors.
• Block importance depends on the calibration data and task. In this small model, data-driven methods converge toward the final blocks, but the calibration data shapes the selection at the margins, enough to produce measurable differences in task performance and an effect that grows with model size.
• Research papers "Shortened LLaMA" (Kim et al., 2024) and "ShortGPT" (Men et al., 2024) validate depth pruning as superior to width pruning for latency reduction and confirm that cosine-based importance metrics are effective across models up to 70B parameters.

FAQ