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_dtypetorch.float32, device_mapcpu ) print(原始模型 config num_attention_heads:, model.config.num_attention_heads) # 每层剪掉2个注意力头Qwen2‑0.5B 原12head →保留10head prune_model prune_attention_heads(model, prune_num_heads2) # 测试推理 prompt 讲一个简短小故事 inputs tokenizer(prompt, return_tensorspt) outputs prune_model.generate( **inputs, max_new_tokens100, do_sampleFalse ) print(\n剪枝后输出) print(tokenizer.decode(outputs[0], skip_special_tokensTrue)) # 保存剪后模型可以后续微调 prune_model.save_pretrained(./qwen2‑0.5b‑pruned‑head) tokenizer.save_pretrained(./qwen2‑0.5b‑pruned‑head)4.剪枝代码优化方向实际需要根据每个head重要性进行剪枝思路跑一部分校准数据统计每个注意力头的平均注意力熵 / 梯度排序移除分数最差的 head。