LLM-Pruner

图1：LLM-Pruner 项目 Logo

背景：当大模型变得臃肿，我们该怎么办？

想象这样一个场景：你的手机里装了一个能力超强但体积上百 GB 的 AI 助手，每次提问都要等上十几秒，还动不动因为内存不足闪退。这是 2023 年大语言模型（LLM）领域面临的真实困境——以 LLaMA、ChatGLM 为代表的大模型虽然能力惊人，但巨大的参数量（7B~65B）让它们在消费级硬件上的部署变得几乎不可能。

如何把大模型压缩到可以本地运行的规模，同时尽量保留它的原有能力？新加坡国立大学（NUS）的研究团队提出了一个优雅的方案——结构化剪枝（Structural Pruning），并将其封装为开源工具 LLM-Pruner，发表在 NeurIPS 2023。

图2：LLM-Pruner 核心原理——对 LLaMA 模型进行结构化剪枝

一、项目起源与学术背景

LLM-Pruner 由新加坡国立大学的研究人员 Xinyin Ma、Gongfan Fang 和 Xinchao Wang 于 2023 年 5 月提出，对应论文《On the Structural Pruning of Large Language Models》，已在 arXiv 公开（arXiv:2305.11627）。项目由 horseee 独立维护和开发，当前在 GitHub 获得 1130 颗星、135 个 Fork，被 NeurIPS 2023 会议接收。

该项目属于模型压缩（Model Compression）领域的一条核心技术路线——结构化剪枝。传统的非结构化剪枝（随机丢弃权重）虽然压缩率高，但生成的稀疏矩阵需要专用硬件才能加速，实际部署收益有限。结构化剪枝则以注意力头、前馈网络层、Embedding 通道等结构单元为最小粒度进行裁剪，剪枝后模型仍是稠密矩阵，可在普通 GPU 上直接加速。

二、核心技术原理

LLM-Pruner 的技术栈建立在三大核心组件之上：

2.1 依赖图剪枝（Dependency Graph Pruning）

LLM-Pruner 自研了一套 TorchPruning 框架，其核心是依赖图（Dependency Graph）数据结构。PyTorch 的 autograd 机制使得模型各层之间存在复杂的依赖关系——剪掉某一层的注意力头，会连带影响其后的 LayerNorm、MLP 和输出投影层。依赖图算法通过追踪 PyTorch 的计算图（grad_fn），自动建立所有层之间的依赖关系，确保剪枝操作是结构一致的——即剪掉一个通道时，与其相连的所有下游层同步更新维度，避免维度错配。

代码实现中，DependencyGraph 类负责遍历模型的 nn.Module 树，为每个模块节点建立输入/输出连接关系（adjacency list）。每个节点维护一个 dependencies 列表，记录哪些操作会受该节点变化的影响。剪枝时，只需指定要移除的通道索引，框架自动通过依赖图推导出所有需要同步修改的层，批量完成剪枝操作。

2.2 重要性评估（Importance Estimation）

剪枝剪哪些？LLM-Pruner 内置了多种重要性评估策略，最常用的是 Magnitude Importance（幅度重要性）——对权重矩阵求 L2 范数，幅度越小的通道对模型输出贡献越低，优先裁剪。同时也支持基于泰勒展开（Taylor Expansion）的一阶导重要性评估方法，该方法利用损失函数对参数的梯度信息，能更精准地评估每个通道对最终 loss 的影响。

评估在 50k 公开 alpaca 数据集上进行，相比从头训练极大降低了计算开销。

2.3 LoRA 后训练（Post-training）

剪枝后模型能力会有所损失，LLM-Pruner 采用 LoRA（Low-Rank Adaptation） 技术进行快速微调恢复。LoRA 的核心思想是在冻结原有权重矩阵的同时，额外注入可训练的秩分解矩阵（Low-Rank Decomposition），参数量仅为原模型的 0.1%~1%，但足以在少量数据上恢复模型能力。

项目直接集成了 peft 库（Parameter-Efficient Fine-Tuning），支持 LoraConfig、get_peft_model 等标准接口，训练后可通过 merge_and_unload() 将 LoRA 权重与原模型合并，得到可直接部署的稠密剪枝模型。

三、支持模型与效果数据

LLM-Pruner 支持当前主流的开源大语言模型家族：

模型	说明
Llama-3.1	Meta 最新一代，最高达 405B 参数
Llama-3 / Llama-2 / LLaMA	Meta 开源系列，7B~65B
Vicuna	基于 LLaMA 的对话微调模型
BLOOM	BigScience 开源多语言模型
ChatGLM	清华 ChatGLM 系列（6B 等）
Baichuan	百川智能开源系列
TinyLlama	小型化 LLaMA 变体

论文中的实验数据表明，在 LLaMA-7B 上剪枝 20% 的参数后，使用 50k alpaca 数据训练 3 小时，困惑度（PPL）仅轻微上升，但推理速度提升约 1.42 倍，显存占用降低 23%。对 ChatGLM-6B 的实验也展示了类似效果，剪枝后模型在多项下游任务上的精度损失控制在 2% 以内。

图3：LLaMA 剪枝前后对比——压缩比与精度保持

图4：ChatGLM-6B 剪枝效果——各任务精度对比

四、代码架构解析

LLM-Pruner 的代码库组织如下：

LLMPruner/
  torch_pruning/     # 自研结构化剪枝核心引擎
    dependency.py   # 依赖图算法（核心，30KB）
    importance.py   # 重要性评估（Taylor / Magnitude）
    ops.py          # 算子类型映射
    _helpers.py     # 索引映射工具（Flatten / Concat）
    pruner/        # 通用剪枝函数注册表
  pruner/          # 面向各模型的剪枝适配器
    hf_llama_pruner.py
    hf_baichuan_pruner.py
    hf_chatglm_pruner.py
  models/          # 各模型结构定义（patched）
  peft/            # PEFT（LoRA）集成封装
  evaluator/       # PPL 困惑度评估器
  datasets/        # 数据集加载
hf_prune.py       # 入口：剪枝阶段
post_training.py   # 入口：LoRA 后训练
generate.py        # Gradio 演示界面
llama3.py         # Llama-3 专用剪枝脚本

整体采用适配器模式：核心引擎 torch_pruning/ 与具体模型无关，通过 pruner/ 中的模型适配器（如 HFRMSNormPrunner、HFAttentionPrunner）处理不同模型的特殊结构。这种设计使得添加新模型支持只需继承 BasePruningFunc 并实现剪枝接口，无需改动核心逻辑。

关键代码：注意力头剪枝

class HFAttentionPrunner(BasePruningFunc):
    def prune_out_channels(self, layer, idxs):
        assert len(idxs) % layer.num_heads == 0
        for sub_layer in [layer.o_proj]:
            keep_idxs = sorted(set(range(sub_layer.out_features)) - set(idxs))
            sub_layer.out_features -= len(idxs)
            sub_layer.weight = torch.nn.Parameter(sub_layer.weight.data[keep_idxs])
            if sub_layer.bias is not None:
                sub_layer.bias = torch.nn.Parameter(sub_layer.bias.data[keep_idxs])

这段代码展示了对 HuggingFace 格式 Llama 模型中 Attention 输出投影层（o_proj）的结构化剪枝——按注意力头粒度对齐移除通道索引，同时更新权重矩阵和偏置向量，确保下游层的输入维度匹配。

五、部署方式与上手体验

LLM-Pruner 支持两种部署路径：

5.1 直接安装（推荐有 GPU 的开发者）

git clone https://github.com/horseee/LLM-Pruner.git
cd LLM-Pruner && pip install -r requirements.txt
bash scripts/llama_prune.sh
python post_training.py --model ... --data alpaca --epochs 3

硬件要求方面，7B 参数模型剪枝需要约 16GB 显存（V100/A100），后训练阶段同样需要 GPU 加速。官方建议使用 RTX 3090 或更高规格的消费级/专业 GPU。

5.2 Docker 一键部署

项目提供了 Dockerfile（多阶段构建）和 docker-compose.yml，可快速构建包含 PyTorch + CUDA + 基础依赖的运行环境，降低环境配置的复杂度。但需注意：LLM 权重文件（通常 15GB~130GB）需单独从 HuggingFace 下载，Docker 镜像本身不包含模型权重。

图5：LLM-Pruner 部署架构——Docker + Gradio 演示界面

5.3 Gradio 演示（体验版）

剪枝完成后，可通过 generate.py 启动一个 Gradio Web 界面，在浏览器中与剪枝后的模型对话。该脚本封装了模型加载、LoRA 权重合并和 transformers 的 GenerationConfig，是一个开箱即用的交互演示。

六、局限性与使用注意事项

尽管 LLM-Pruner 在模型压缩领域有重要的学术和实践价值，但使用时需注意以下局限：

1. 量化兼容性：当前版本不直接支持与 INT4/INT8 量化联合使用。量化通常在剪枝后额外执行，联合优化策略（如 QLoRA）需要自行组合。

2. 特定任务场景：剪枝后的模型在通用 NLP 任务上表现稳定，但在极 specialized 的下游任务（如特定领域知识问答）上，能力下降可能更明显，需要更长的后训练。

3. 训练数据依赖：后训练默认使用 alpaca 数据集。如果要在中文场景使用，建议替换为中文指令微调数据集（如 Belle、Firefly），否则模型可能在中英文混用场景下效果不佳。

4. 非 HuggingFace 模型：项目当前主要针对 HuggingFace transformers 格式的模型优化，对纯 PyTorch 模型的支持需要额外适配工作。

七、行业意义与未来方向

LLM-Pruner 代表了大模型压缩领域的一个重要分支——结构化剪枝。随着 Llama-3、Qwen、DeepSeek 等开源大模型的持续涌现，如何在消费级硬件上高效运行这些模型，成为 AI 落地部署的核心挑战之一。

从技术趋势看，剪枝+量化+知识蒸馏的联合优化（剪-quant-蒸三件套）是当前最主流的 LLM 压缩路线。LLM-Pruner 填补了剪枝环节的开源工具空白，与 llama.cpp（量化）、Axolotl（微调）等工具共同构成了开源 LLM 部署工具链的重要一环。

作者在 README 中提到正在开发自动化结构化剪枝功能，未来用户只需指定目标模型名称和压缩比例，工具链将自动完成发现→剪枝→后训练全流程，进一步降低使用门槛。

图6：ChatGLM-6B 详细实验数据——各压缩比例下的任务精度

---

总结：LLM-Pruner 是大模型结构化剪枝领域难得的工程级开源实现，将学术论文中的方法论转化为可跑、可改、可扩展的完整工具链。对 AI 研究者和工程师来说，它是探索 LLM 高效部署的一把利器。

项目链接：https://github.com/horseee/LLM-Pruner | 论文：arXiv:2305.11627 | 会议：NeurIPS 2023