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。