
AMCT 结构化剪枝实战基于 amct_pytorch.pruning 对 FFN 中间维 / CNN 通道 / MoE 专家进行原地压缩【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct本文是 CANN AMCT 中amct_pytorch.pruning模块的结构化剪枝使用指南。该模块面向已实例化并加载权重的torch.nn.Module支持对稠密 FFN 中间维、CNN 卷积通道与 MoE 路由专家三种结构进行结构化剪枝模型被原地修改剪枝结果可同步回写model.config从而支持直接save_pretrained与重载。读完本文你将掌握容差驱动的自动剪枝搜索、固定剪枝率的精确控制、CNN 重建式剪枝、MoE 专家方差菜单择优、恢复菜单、跨层敏感度分配以及剪枝与量化串联部署的完整流程。核心设计一次前向、四步完成结构化剪枝结构化剪枝按打分 → 剪枝 → 恢复 →可选量化四步执行。工具接收一个已加载权重的模型和少量校准数据几批前向即可在校准数据上前向一遍为每个可剪结构计算“重要性分数”移除分数最低的部分必要时对保留下来的权重做一次轻量补偿恢复随后可选接续量化。整个过程原地修改模型并同步改写config中的维度如intermediate_size/num_experts因此剪枝后的模型可直接save/load无需额外转换步骤。本库从不下载模型模型实例化由调用方负责也不保留任何进程级全局状态统计信息统一通过可选的reportPruneReport()出口返回。入口函数定义于 api.py签名如下搜索模式与固定模式的完整参数amct.prune( model, # torch.nn.Module原地修改 configNone, # dict / PruneConfig / None dataNone, # 校准数据 toleranceNone, # 传入则进入容差搜索 evaluatorNone, # callable(model)-float或暴露 .evaluate(model) 的对象 eval_iterations1, eval_dataNone, ratio_gridNone, # 候选剪枝率网格 size_budgetNone, # 传入则进入 size budget 搜索与 tolerance 互斥 finetune_fnNone, quant_fnNone, reportNone, # PruneReport() 出口 )从源码可见 api.py 对tolerance与size_budget同时传入会直接抛ValueError二者是互斥的搜索模式。三个可剪枝的域能剪什么、剪完之后什么保持不变取决于“域”。工具只剪生产者→消费者接口能对得上的维度对不确定的一律保守跳过保证张量形状始终自洽稠密 FFN—— 只收窄中间维gate/up_proj的输出通道和down_proj的输入通道一起变小隐藏维 / 残差宽度保持不动注意力 q/k/v/o 投影自动排除。CNN 通道—— 沿“生产者卷积输出通道 →可选 BatchNorm→ 消费者输入通道”这条链一起改尺寸喂入残差add的卷积、ConcatInception消费方、分组 / 深度可分离卷积都不剪。MoE 专家—— 整块移除被路由的专家并同步收缩路由器共享始终激活专家保留隐藏 in/out 不变。域检测的实现分别位于 domains/dense.py、domains/cnn.py 与 domains/moe.py。快速上手容差驱动的自动剪枝只需指定一个可接受的精度损失tolerance工具即在ratio_grid上二分查找每试一个剪枝率r试剪一遍按方法以掩码在原模型上进行或在副本上、可选微调后测量精度下降量下降在容差内则搜索转向更大的r超出则退回更小的r。最终应用满足容差的最大剪枝率。探测某一剪枝率时的任何剪枝 / 前向失败均视为不可接受不抛异常因此在非单调情形下搜索倾向于少剪。搜索核心实现在 accuracy_based_auto_prune.py其中DEFAULT_RATIO_GRID (0.1, 0.2, ..., 0.8)、DEFAULT_TOLERANCE 0.02见该文件 L38-L40。import amct_pytorch as amct amct.prune(model, datacalib, tolerance0.02) # 在容差内搜索并应用最大剪枝率原地生效⚠️搜索期间的内存开销取决于所用方法。只做“选哪些保留”的方法low_variance、variance_channel、activation_count、mass_variance以掩码量测候选剪枝率把剪枝以掩码形式作用于原模型、量测后原样还原设备上始终只有一份模型。改写权重的方法reconstruct的最小二乘、output_merge的专家合并无法用掩码表达仍在整模型副本上试剪峰值内存约为模型的两倍以上传入finetune_fn/quant_fn时同理回调会修改模型敏感度分配allocation.strategysensitivity逐层测敏感度时也仍在整模型副本上进行。固定剪枝率模式从不复制。当提供finetune_fn时每个候选剪枝率在评估前都会先微调因此被选中的剪枝率是剪枝 恢复后仍可用的那个from amct_pytorch.pruning import prune_finetune amct.prune(model, cfg, datacalib, tolerance0.05, evaluatortest_acc, finetune_fnlambda m: prune_finetune(m, train_data, steps300))默认保真度指标无需标签质量 在校准数据上与原始模型 top-1 预测的一致率。也可改为传入evaluatorcallable(model)-float或任何暴露.evaluate(model) - float的对象。prune 的参数参数是否必填默认值说明model是-待剪枝的torch.nn.Module原地修改data视情况None校准数据基于方差的方法、默认评估集需要tolerance否None可接受精度损失的上界与evaluator同量纲传入即进入容差搜索evaluator否top-1 保真度callable(model)-float或暴露.evaluate(model) - float的对象eval_data否回退到data默认保真度指标使用的评估批ratio_grid否0.1..0.8候选剪枝率升序report否None传入PruneReport()作为出口以取回统计信息注意力投影默认被跳过由skip_layers配置控制见下文“进阶配置”。evaluator/eval_data/ratio_grid/finetune_fn/quant_fn仅在搜索模式下生效固定剪枝率下传入会报ValueError不会被静默忽略见 api.py 的_reject_search_args。固定剪枝率精确控制每个域from amct_pytorch.pruning import PruneReport cfg {methods: {dense: {name: low_variance, kwargs: {prune_ratio: 0.5}}}} rep PruneReport() amct.prune(model, cfg, datacalib, reportrep) print(rep.as_dict())methods里列出哪些域就只剪哪些域。未列出的域一律固定为prune_ratio: 0.0不会按默认率被顺带剪掉固定剪枝率、容差搜索、size budget、菜单择优、诊断都遵循这一条。configNone一个域都不指定时才回到三域默认值 cnn 0.30 / dense 0.50 / moe 0.50。这一语义在 config.py 的resolved_methods()中实现未命名任何域时使用默认方法表cnnvariance_channel0.30、denselow_variance0.50、moeactivation_count0.50一旦命名了任一域其余域被显式置为prune_ratio: 0.0。kwargs只接受该方法实际读取的键。写错的键例如把prune_ratio写成prune_rate会抛ValueError并列出该方法可用的键而不是被静默忽略、回落到默认剪枝率。各内置方法声明的合法键accepted_kwargs定义于 prune_op/base.py如下方法所属域接受的 kwargslow_variancedenseprune_ratiovariance_channelcnnprune_ratioreconstructcnnprune_ratio,ridge,recoveryreconstructdenseprune_ratio、ridge、recovery以及量化感知相关的quant_cfg见 dense_reconstruct.pyactivation_countmoeprune_ratio,top_kmass_variancemoeprune_ratio,top_k,boundary,variance_scoreoutput_mergemoekeep_ratio默认 0.50另有prune_ratio别名 1 - keep_ratio、selectorcalib_nll默认 /none校验逻辑位于 config.pyprune_ratio必须为[0.0, 1.0)内的数值未知键会报错并列出可接受键。PruneReport 结构prune(..., reportrep)就地填充rep无进程级全局状态通过rep.as_dict()读取字段定义见 report.py字段含义backend识别到的模型后端huggingface/modelscope/pretrained-module/torchparams_before/params_after剪枝前 / 后的参数量prunable_fraction位于可剪结构中的参数占比仅size_budget模式下填充per_layer_sparsity{module_path: 稀疏率}—— 各被剪层移除的比例warnings非致命提示跳过的融合专家、非均匀宽度等budget_unreachable若size_budget目标无法达成则为Trueallocation_choice采用的跨层分配策略均匀 / 敏感度events按 (域, 方法, 模块) 记录剪了什么后端识别实现在 compat.py 的detect_backend依据模块所属包transformers/modelscope、是否暴露save_pretrainedconfig判定。CNN 通道剪枝variance_channel按激活方差切分通道同父兄弟卷积启发式。reconstruct先移除低显著度通道再通过 im2col 最小二乘重建消费方卷积权重以保持输出无需微调from amct_pytorch.pruning import CNN_RECONSTRUCT_PRUNE_CFG amct.prune(model, CNN_RECONSTRUCT_PRUNE_CFG, datacalib) # custom: {methods: {cnn: {name: reconstruct, # kwargs: {prune_ratio: 0.3, ridge: 1e-2}}}}输出头最后一个 Linear及匹配skip_layers的层会被自动忽略。校准必须提供足够多的行 keep x k x k以求解最小二乘当方程欠定或消费方不是 Conv2d 时会发出警告并改用朴素切片。在残差ResNet块中内部 conv-conv 耦合会被剪枝馈入残差add的卷积会被自动排除。ConcatInception消费方以及分组/深度可分离卷积不剪枝。ridge参数为最小二乘加入岭正则CNN_RECONSTRUCT_PRUNE_CFG中默认prune_ratio0.30见 presets.py。MoE 专家剪枝方差准则菜单进阶不在__all__中—— 通过from amct_pytorch.pruning import ...导入。将MOE_VARIANCE_MENU_CFG传给prune即可运行一次校准、按菜单中的每个准则各剪枝一份副本、在小验证集上挑出最佳并应用胜者。prune检测到配置带menu就自动改走菜单择优无需再传tolerance。“菜单”类配置的套路是一次校准、多方案择优共用同一遍校准把菜单里每个候选各剪出一份副本在一个独立的小验证集eval_data上实测谁严格赢过安全回退项才用谁打平或没赢就保留回退项。如此某个在校准集上好看、但在验证集上并不稳定的方案便不会被误选。MoE 方差菜单里各候选即不同的重要性打分准则mass/cond_var/peak/cvxpeak其中mass是那个安全回退项。菜单实现在 presets.py择优判定“严格胜出”的逻辑见 accuracy_based_auto_prune.py 中的_menu_select。from amct_pytorch.pruning import MOE_VARIANCE_MENU_CFG amct.prune( model, MOE_VARIANCE_MENU_CFG, datacalib, eval_dataval, evaluatorev) # keep val ! calib菜单massactivation_count回退/cond_var/peak/cvxpeak。方差准则仅在验证集上严格胜出时才替换mass否则回退到mass。各准则的分数是mass_variance的 kwargvariance_score取值 {cond默认,peak,cvxpeak}深度切分boundary层 boundary 用 mass层 boundary 用 variance默认10-1 全用 varianceauto 逐层list[int] 指定这些层。MOE_VARIANCE_MENU_CFG默认boundary-1传入带boundaryK的自定义菜单可搜索深度切分。恢复菜单稠密 FFN / CNN 通道剪枝进阶不在__all__中—— 通过from amct_pytorch.pruning import ...导入。将DENSE_RECOVERY_MENU_CFG传入以搜索剪枝后的恢复菜单。恢复菜单沿用完全相同的“择优菜单”机制只是候选换成了剪枝后的补偿方式——同样是一次校准、在独立验证集上择优、只有严格胜出才替换回退项from amct_pytorch.pruning import DENSE_RECOVERY_MENU_CFG amct.prune( model, DENSE_RECOVERY_MENU_CFG, datacalib, eval_dataval, evaluatorev)CNN 通道剪枝用同一套菜单换成CNN_RECOVERY_MENU_CFG菜单模式只剪菜单所在的域from amct_pytorch.pruning import CNN_RECOVERY_MENU_CFG amct.prune( model, CNN_RECOVERY_MENU_CFG, datacalib, eval_dataval, evaluatorev)菜单none朴素切片回退/bias均值折叠进消费方 bias/ls最小二乘重建。由reconstruct的 kwargrecovery支撑取值 {ls默认,bias,none}。菜单定义见 presets.py。仅在验证集上严格胜出时才替换回退项请使用能代表部署场景的eval_data。校准充足时倾向ls校准量 中间维宽度会使ls欠定此时bias更稳健。bias折叠可能在偏移的测试集上过拟合校准均值为保留测试侧下限可从cfg[methods][dense][menu]中去掉bias。CNN 的 im2col 行数充足因此ls在那里不会欠定。适用范围无微调场景若剪枝后进行微调朴素切片可能追平。进阶配置import amct_pytorch as amct from amct_pytorch.pruning import ( SENSITIVITY_ALLOC_PRUNE_CFG, MOE_OUTPUT_MERGE_PRUNE_CFG, ) # cross-layer sensitivity allocation (default uniform) amct.prune(model, SENSITIVITY_ALLOC_PRUNE_CFG, datacalib) # config[allocation] {strategy: sensitivity, ref_ratio: 0.5, # min_ratio: 0.05, max_ratio: 0.9, guard: calib_nll} # 注意敏感度分配会让各层宽度不一致而 config 中相应字段只有单个标量无法描述这种结构 # 因此计数字段不会被写入保存后也无法用原 from_pretrained 还原需要自行保存每层尺寸并在 # 加载时重建结构。需要标准保存/加载时请让各层使用同一剪枝率。详见「剪枝后的 config 同步」。 # MoE output-space expert merge (LS residual pairing drop fallback) amct.prune(model, MOE_OUTPUT_MERGE_PRUNE_CFG, datacalib) # moe kwargs: keep_ratio0.5, selectorcalib_nll|none # skip_layers: 名称匹配这些子串的层不剪。搜索模式tolerance / size_budget / menu会自动 # 并入注意力投影 self_attn/attention/attn固定剪枝率模式需要自行列出。 # quant_cfg 只被 reconstruct 方法读取low_variance 会忽略它 cfg {methods: {dense: {name: reconstruct, kwargs: {prune_ratio: 0.5}}}} cfg[skip_layers] [self_attn, lm_head, shared_expert] # quantization-aware saliency cfg[methods][dense][kwargs][quant_cfg] { weights_cfg: {strategy: channel, symmetric: True, dtype: int8}} # size budget: keep 70% of the parameters amct.prune(model, cfg, datacalib, size_budget0.7)上述预设常量SENSITIVITY_ALLOC_PRUNE_CFG、MOE_OUTPUT_MERGE_PRUNE_CFG等定义于 presets.py并统一从 pruning/init.py 导出。关于敏感度分配allocation.strategysensitivity时跨层剪枝率由各层敏感度决定均匀分配为默认敏感度测量、water_fill_ratios水位线分配等实现在 allocation.py。allocation的合法性校验见 config.pystrategy仅允许uniform/sensitivityguard仅允许calib_nll/none各 ratio 均须在[0, 1)且min_ratio max_ratio。显存受限时的手动循环使用reconstruct/output_merge搜索或传入finetune_fn/quant_fn时仍会复制整个模型。若设备显存不足以容纳两份可把干净权重留在主机内存自行在固定剪枝率上循环——固定剪枝率不复制设备上始终只有一份工作模型import copy import torch import amct_pytorch as amct pristine model.cpu() # 基线留在主机内存不参与计算 best None for prune_ratio in (0.3, 0.4, 0.5): trial copy.deepcopy(pristine).to(npu:0) # 设备上只有这一份 cfg {methods: {dense: {name: low_variance, kwargs: {prune_ratio: prune_ratio}}}} amct.prune(trial, cfg, datacalib) # 固定剪枝率不复制 if my_evaluator(trial) threshold: # 自行评估并保留满足要求的最大剪枝率 best prune_ratio del trial torch.npu.empty_cache()注意PruneConfig.copy_model字段默认True见 config.py在固定模式下会被置为False见 api.py即固定剪枝率始终原地操作不复制。与量化结合使用先剪枝再量化最后转换剪枝会改变张量形状因此量化校准只有在剪枝后的权重上才正确import amct_pytorch as amct amct.prune(model, datacalib, tolerance0.02) # 1) structured pruning amct.quantize(model, quant_cfg) # 2) quantize on the pruned model amct.convert(model) # 3) convert the deployment model恢复与量化回调finetune_fn / quant_fn两者都是可选回调在容差搜索中对每个剪枝副本评估前应用因此被选中的剪枝率反映的是恢复后 / 量化后的精度finetune_fn(model)—— 对剪枝副本做一次轻量恢复如几步优化。仓库提供了开箱即用的 prune_finetune默认steps200、lr1e-4、weight_decay0.0、warmup0、grad_clip1.0使用 AdamW默认损失对 dict 批取 causal-LM loss、对 (x, y) 元组批取交叉熵也可传入自定义loss_fn(model, batch)。quant_fn(model)—— 对剪枝副本施加量化使搜索计入剪枝量化的合并损失只会让搜索剪得更保守不会更激进。from amct_pytorch.pruning import prune_finetune # 仅恢复每个候选评估前先微调 amct.prune(model, cfg, datacalib, tolerance0.05, evaluatorev, finetune_fnlambda m: prune_finetune(m, train_data, steps300)) # 剪枝 量化共用一个容差搜索看到的是量化后的精度 amct.prune(model, cfg, datacalib, tolerance0.03, evaluatorev, quant_fnlambda m: amct.quantize(m, quant_cfg))evaluator 协议剪枝的evaluator只要求一个协议对象暴露evaluate(model)并返回精度指标越高越好。不依赖任何基类因此同一个评估器也可直接传给量化的accuracy_based_auto_calibrationclass MyEvaluator: def evaluate(self, model): return my_top1_accuracy(model) ev MyEvaluator() amct.prune(model, datacalib, tolerance0.02, evaluatorev) amct.accuracy_based_auto_calibration(model, ev, quant_cfg, ...)内置的amct.ModelEvaluator是量化校准的数据馈送器其evaluate(model, iterations)返回 None不能用于精度搜索。精度搜索请传入返回指标的评估器。剪枝侧同时接受evaluate(model)和evaluate(model, iterations)。支持的可剪枝结构只有生产方↔消费方接口可被验证的维度才会被剪枝其余跳过。稠密 FFN—— 仅剪枝中间维gate/up.out_featuresdown.in_features隐藏/残差宽度保持不变。注意力投影q/k/v/o被自动排除。CNN 通道—— 生产方Conv2dgroups1- 可选BatchNorm2d- 通道匹配的消费方Conv2d/Linear一并改尺寸。馈入残差add的卷积自动排除Concat 消费方以及分组/深度可分离卷积不剪枝。MoE 专家—— 整个被路由的专家被移除隐藏 in/out 不变。共享始终激活专家被排除。域是否支持备注dense是三 Lineargate/up/down_proj、融合gate_up_projPhi-3/GLM-4、相邻两 Linear/Conv1D含 Bloom 风格。Llama/Qwen2/Mistral/Qwen3 无需手动skip_layers即可剪枝。cnn是variance_channel启发式切片reconstructim2col 最小二乘。同层检测无法识别 ResNet 风格的独立块。moe是传统nn.ModuleListnn.Linear门控、融合批处理专家MixtralExperts/Qwen3MoeExperts*TopKRouter、分组路由器n_group/topk_group、共享专家 sigmoid 路由noaux_tc、兄弟双张量融合专家GraniteMoE、嵌套路由器 biasErnie4.5。机械可剪枝性在多种常见架构家族Llama/Qwen/Mixtral/GLM/Phi/GPT/GraniteMoE/Ernie4.5 等的 2 层微模型上有覆盖测试见 test_pruning_real_hf_models.py 与 test_pruning_auto_prune.py。剪枝后的 config 同步剪枝结束后会把新尺寸写回model.config使save_pretrained保存的配置与权重一致。该逻辑实现在 compat.py 的patch_common_config稠密 FFN 中间维 -intermediate_size/ffn_hidden_size/n_inner属性名集合见 compat.py。专家数 -num_local_experts/num_experts/n_routed_experts/n_expertscompat.py。每 token 专家数 -num_experts_per_tok/moe_top_k/top_k/num_selected_experts会被下调到不超过剩余专家数否则路由器会选到已删除的专家首次前向即报错compat.py。多模态包装模型会通过白名单text_config/llm_config/language_config定位文本子配置写入避免误改视觉塔的intermediate_sizecompat.py。各层专家数不一致时单个标量无法描述非均匀剪枝因此num_experts一类的计数字段不会被写入只打印告警——写入任一层的数值都会与其余层不符。此时top_k仍会被下调到最薄一层的专家数保证所有层都能前向。这类模型直接save_pretrained后无法用原from_pretrained还原需要自行保存每层尺寸并在加载时重建结构。若要避免该情况请让各层使用同一剪枝率prune_ratio不要使用会产生非均匀结果的敏感度分配SENSITIVITY_ALLOC_PRUNE_CFG。稠密 FFN 各层宽度不一致时同理。故障排查与诊断入口若不确定模型结构能否被识别可先运行 prune_diagnose 做一次不修改原模型的预检它会统计三个域各自识别到的可剪目标数、在副本上试剪一次固定剪枝率并检查前向是否通过、再以applyFalse干跑一次容差二分搜索返回DiagnosisReport含prune_works、search_chosen_ratio与备注便于在正式剪枝前定位结构不匹配或数据缺失问题。更多示例仓库的剪枝示例CNN / 稠密 / MoE 三种场景的完整脚本位于 examples/algorithms/pruning/src包含基于激活次数预算与容差两种约束的 Qwen3 6.35B MoE 剪枝样例可作为从本文 API 到真实模型应用的衔接参考。【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考