模型压缩有三板斧量化已讲、剪枝Pruning、蒸馏Distillation。剪枝去掉不重要的权重蒸馏用小模型学大模型的输出。在昇腾NPU上剪枝和蒸馏可以联合使用把 7B 模型压到 3B 精度损失 ❤️%。剪枝Pruning去掉权重矩阵中绝对值最小的连接结构化剪枝去掉整个神经元或 attention head。fromtorch_npu.contribimportprune modelAutoModelForCausalLM.from_pretrained(meta-llama/Llama-2-7b-hf,torch_dtypetorch.bfloat16,device_mapnpu:0,)# 结构化剪枝去掉 25% 的 attention headprunerprune.StructuredPruner(model,pruning_ratio0.25,# 剪掉 25% 的 headcriterionl1_norm,# 按 L1 范数排序小的去掉blocks[self_attn],# 只剪 attention)# 剪枝pruned_modelpruner.prune()# 微调恢复精度train(pruned_model,dataloader,epochs3)# 保存剪枝后的模型结构变了需要特殊格式pruner.save(llama2-7b-pruned-25pct.pt)剪枝后的模型结构变了attention head 数减少不能直接用标准 LLM 加载。需要 ATB 支持动态 head 数——目前 CANN 8.5 还不支持需要等后续版本。蒸馏Distillation大模型Teacher指导小模型Student训练importtorchimporttorch_npufromtransformersimportAutoModelForCausalLM# Teacher: Llama2-13BteacherAutoModelForCausalLM.from_pretrained(meta-llama/Llama-2-13b-hf,torch_dtypetorch.bfloat16,device_mapnpu:0,1,# 13B 需要 2 卡)# Student: Llama2-7BstudentAutoModelForCausalLM.from_pretrained(meta-llama/Llama-2-7b-hf,torch_dtypetorch.bfloat16,device_mapnpu:0,)# 蒸馏训练optimizertorch.optim.AdamW(student.parameters(),lr1e-5)forbatchindataloader:# Teacher 前向不计算梯度withtorch.no_grad():teacher_logitsteacher(batch[input_ids]).logits# Student 前向student_logitsstudent(batch[input_ids]).logits# 蒸馏损失KL 散度T2.0# 温度losstorch.nn.functional.kl_div(torch.nn.functional.log_softmax(student_logits/T,dim-1),torch.nn.functional.softmax(teacher_logits/T,dim-1),reductionbatchmean,)*(T*T)loss.backward()optimizer.step()optimizer.zero_grad()蒸馏 剪枝联合先蒸馏训练一个较大的 Student如 7B再剪枝到更小如 4B最后微调恢复精度步骤 1蒸馏13B → 7B训练 3 轮 步骤 2结构化剪枝7B → 4B去掉 40% 的 FFN 神经元 步骤 3用蒸馏损失继续微调 1 轮恢复精度Llama2-13B → 7B → 4B 的精度损失阶段WNLI (准确率)GSM8K (准确率)模型大小Teacher (13B)78.5%56.2%26GBStudent (7B蒸馏)76.8% (-1.7%)54.1% (-2.1%)14GBPruned (4B)74.2% (-4.3%)51.3% (-4.9%)8GBPruned 微调75.9% (-2.6%)53.7% (-2.5%)8GB最终 4B 模型精度损失 2.6%大小减到原来的 30%。昇腾NPU上的加速蒸馏训练需要同时跑 Teacher 和 Student前向 反向。显存需求Teacher (13B, fp16): 26GB冻结不存梯度 Student (7B, fp16): 14GB 28GB优化器 42GB 激活: 10GB 总计: 26 42 10 78GB → 需要 2 卡128GB13B Teacher 放在 2 卡 TP7B Student 放在 1 卡剩余 1 卡做梯度累积。跟量化的配合蒸馏/剪枝后的模型再做量化进一步压缩原始 7B: 14GB (fp16) → 蒸馏到 4B: 8GB (fp16) → W4A16 量化: 2.5GB (int4)2.5GB 的模型可以放到 Jetson Orin32GB 显存上跑推理延迟 50ms/token。剪枝和蒸馏是模型压缩的高阶技巧。剪枝在 CANN 8.5 上支持有限动态 head 数待支持蒸馏已经可以完整跑通。两者联合使用7B → 4B 精度损失 ❤️%。仓库在这里https://atomgit.com/cann/torch_npu
