18.剪枝

1.剪枝的定义

把神经网络里不重要的权重 / 神经元删掉，在尽量少掉精度前提下，减小模型体积、降低显存、加速推理，常用于端侧部署。

2.剪枝的分类

• 非结构化剪枝 -- 减少权重

NPU 无法加速，且 RKNN 不支持稀疏推理，RK3588上基本无意义（RK3588不建议做）。

权重就是神经网络里那些被训练出来、用来“加权计算”的数字；

• 结构化剪枝 -- 删除Attention Head / 整个 Layer/某些通道

改变矩阵形状，需重新保存 HuggingFace 模型，再由 RKLLM-Toolkit 量化转换

这里的 FNN 通常指 Feedforward Neural Network（前馈神经网络）。 **FFN是 Transformer 中负责“特征变换和非线性增强”的部分，它在每个 token 上独立计算，不做注意力，也不跨 token 交互。

可以理解为：Attention 负责“看哪里重要”，FFN 负责“把信息加工得更丰富”。

如果删减通道，要按通道重要性排序再删，而不是随机删

3.剪枝示例

结构化剪枝代码

import torch from transformers import AutoModelForCausalLM, AutoTokenizer def prune_attention_heads(model, prune_num_heads: int): """ 结构化剪枝：剪掉每一层末尾 prune_num_heads 个注意力头 :param model: causal lm model :param prune_num_heads: 每层要剪掉多少个head :return: pruned model """ config = model.config num_heads = config.num_attention_heads hidden_size = config.hidden_size head_dim = hidden_size // num_heads assert prune_num_heads < num_heads, "剪枝head数不能大于总head" keep_num_heads = num_heads - prune_num_heads for layer_idx, layer in enumerate(model.model.layers): attn = layer.self_attn # ========== Q K V 权重结构化裁剪：保留前 keep_num_heads 个head ========= # weight shape: [hidden_size, hidden_size] q_proj = attn.q_proj.weight k_proj = attn.k_proj.weight v_proj = attn.v_proj.weight # 只保留前 keep_num_heads head 对应的权重 keep_dim = keep_num_heads * head_dim attn.q_proj.weight = torch.nn.Parameter(q_proj.weight[:keep_dim, :].clone()) attn.k_proj.weight = torch.nn.Parameter(k_proj.weight[:keep_dim, :].clone()) attn.v_proj.weight = torch.nn.Parameter(v_proj.weight[:keep_dim, :].clone()) # bias同理 if attn.q_proj.bias is not None: attn.q_proj.bias = torch.nn.Parameter(attn.q_proj.bias[:keep_dim].clone()) attn.k_proj.bias = torch.nn.Parameter(attn.k_proj.bias[:keep_dim].clone()) attn.v_proj.bias = torch.nn.Parameter(attn.v_proj.bias[:keep_dim].clone()) # ========== output proj：输入维度改变 ========== # o_proj: [hidden_size, keep_dim] attn.o_proj.weight = torch.nn.Parameter(attn.o_proj.weight[:, :keep_dim].clone()) if attn.o_proj.bias is not None: pass # 更新模型config，必须！推理时会读取这个配置 config.num_attention_heads = keep_num_heads if hasattr(config, "num_key_value_heads"): # GQA模型需要同步修改KV头，这里简单处理和num_attention_heads保持一致 config.num_key_value_heads = keep_num_heads print(f"剪枝完成：原heads={num_heads}, keep heads={keep_num_heads}, prune={prune_num_heads}") return model if __name__ == "__main__": model_name = "Qwen/Qwen2-0.5B-Instruct" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.float32, device_map="cpu" ) print("原始模型 config num_attention_heads:", model.config.num_attention_heads) # 每层剪掉2个注意力头，Qwen2‑0.5B 原12head →保留10head prune_model = prune_attention_heads(model, prune_num_heads=2) # 测试推理 prompt = "讲一个简短小故事" inputs = tokenizer(prompt, return_tensors="pt") outputs = prune_model.generate( **inputs, max_new_tokens=100, do_sample=False ) print("\n剪枝后输出：") print(tokenizer.decode(outputs[0], skip_special_tokens=True)) # 保存剪后模型，可以后续微调 prune_model.save_pretrained("./qwen2‑0.5b‑pruned‑head") tokenizer.save_pretrained("./qwen2‑0.5b‑pruned‑head")

4.剪枝代码优化方向

实际需要根据每个head重要性进行剪枝

思路：跑一部分校准数据，统计每个注意力头的平均注意力熵 / 梯度，排序，移除分数最差的 head。