TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation

TL;DR

TraceNAS is a training-free neural architecture search method that efficiently identifies optimal structured pruning configurations for large language models by using a zero-shot proxy to maintain loss landscape alignment, reducing computational costs significantly.

Contribution

It introduces a novel zero-shot NAS framework for LLM pruning that captures global dependencies without training, enabling fast and effective model compression.

Findings

Achieves 10× reduction in GPU hours compared to training-aware methods.

Maintains competitive performance on benchmarks with significantly less computation.

Effective in pruning Llama and Qwen models while preserving accuracy.

Abstract

Structured pruning is essential for efficient deployment of Large Language Models (LLMs). The varying sensitivity of LLM sub-blocks to pruning necessitates the identification of optimal non-uniformly pruned models. Existing methods evaluate the importance of layers, attention heads, or weight channels in isolation. Such localized focus ignores the complex global structural dependencies that exist across the model. Training-aware structured pruning addresses global dependencies, but its computational cost can be just as expensive as post-pruning training. To alleviate the computational burden of training-aware pruning and capture global structural dependencies, we propose TraceNAS, a training-free Neural Architecture Search (NAS) framework that jointly explores structured pruning of LLM depth and width. TraceNAS identifies pruned models that maintain a high degree of loss landscape alignment with the pretrained model using a scale-invariant zero-shot proxy, effectively selecting models that exhibit maximal performance potential during post-pruning training. TraceNAS is highly efficient, enabling high-fidelity discovery of pruned models on a single GPU in 8.5 hours, yielding a 10$\times$ reduction in GPU-hours compared to training-aware methods. Evaluations on the Llama and Qwen families demonstrate that TraceNAS is competitive with training-aware baselines across commonsense and reasoning benchmarks.