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.
Abstract (translated)
结构化剪枝对于大型语言模型(LLM)的高效部署至关重要。由于LLM子模块对剪枝敏感度不同,因此需要识别最优的非均匀剪枝模型。现有的方法分别评估层、注意力头或权重通道的重要性,这种局部关注忽略了模型内部复杂的全局结构性依赖关系。训练感知结构化剪枝可以解决这一问题,但其计算成本可能与后剪枝训练一样昂贵。为了减轻训练感知剪枝的计算负担并捕获全局结构性依赖关系,我们提出了TraceNAS,这是一种无需训练的神经架构搜索(NAS)框架,它同时探索LLM深度和宽度的结构化剪枝。TraceNAS通过使用比例不变的零样本代理来识别与预训练模型损失景观高度对齐的剪枝模型,有效地选择在后剪枝训练中表现出最大性能潜力的模型。TraceNAS效率极高,在单个GPU上仅需8.5小时即可发现高质量的剪枝模型,相比训练感知方法减少了10倍的GPU时间消耗。在Llama和Qwen家族上的评估表明,TraceNAS在常识和推理基准测试方面与训练感知基线具有竞争力。
URL
https://arxiv.org/abs/2602.02891