CANN-昇腾NPU-LoRA微调-显存只占5%怎么做到的

全量微调 Llama2-7B 需要更新 7B 参数显存开销约 80GB。LoRA 只训练 0.5% 的参数约 35M显存开销约 4GB。在昇腾NPU上 LoRA 微调是性价比最高的方案。LoRA 原理在原始权重 W 旁边加一个低秩矩阵 ΔW A × B原始y Wx LoRAy (W AB)x Wx ABx A: [hidden, r], B: [r, hidden], r hidden Llama2-7B: r16, A[4096,16], B[16,4096] 每层 LoRA 参数: 2 × 4096 × 16 128K 全量参数: 4096 × 4096 16M 参数比: 0.8%训练时 W 冻结只更新 A 和 B。推理时把 ΔW 合并回 W零额外延迟。昇腾NPU上的实现frompeftimportLoraConfig,get_peft_modelimporttorchimporttorch_npu modelAutoModelForCausalLM.from_pretrained(meta-llama/Llama-2-7b-hf,torch_dtypetorch.bfloat16,device_mapnpu:0)lora_configLoraConfig(r16,lora_alpha32,target_modules[q_proj,v_proj,k_proj,o_proj,gate_proj,up_proj,down_proj],lora_dropout0.05,)modelget_peft_model(model,lora_config)model.print_trainable_parameters()# 输出: trainable params: 35M || all params: 6.7B || trainable%: 0.52%显存对比Llama2-7B 全量微调 vs LoRA 微调单卡 Atlas 800I A264GB 显存项目全量微调LoRA 微调模型权重14GB可冻结14GB冻结LoRA 权重-0.2GB梯度14GB全量0.2GB仅 LoRA优化器状态28GB全量 fp320.4GB仅 LoRA激活20GB20GB总计76GB34.8GBLoRA 微调的显存不到全量微调的一半。单卡 64GB 就能跑不需要 8 卡 TP。训练速度配置吞吐 (tokens/s)训练时间1M tokens全量微调8 卡 TP1,8009 分钟LoRA单卡1,20014 分钟LoRA 单卡比全量 8 卡慢 55%。但如果你只有一张卡LoRA 是唯一选择。推理时合并 LoRA# 方法 1动态合并每次推理时计算 Wx ABx# 优点可以热切换不同 LoRA# 缺点多一次矩阵乘法约 3-5ms 额外延迟# 方法 2预合并训练完成后 W W AB删除 A 和 B# 优点推理零开销# 缺点不能切换 LoRAmodelmodel.merge_and_unload()# 预合并预合并后的模型跟原始模型完全一样——ATB 的推理加速、FlashAttention、量化全部适用。多 LoRA 推理同一个基座模型挂多个 LoRA不同请求用不同 LoRAfromatbimportLLM,MultiLoraConfig modelLLM(meta-llama/Llama-2-7b-hf,devicenpu:0,multi_loraMultiLoraConfig(lora_dirs[lora_chat,lora_code,lora_translate],max_loras3,))# 不同请求用不同 LoRAresultsmodel.generate([(Hello,lora_chat),(def fib(n):,lora_code),(Translate:,lora_translate),])ATB 的 Multi-LoRA 把多个 LoRA 的 ΔW 打包成 Batch GEMM一次计算所有 LoRA 的增量。这比逐个 LoRA 计算高效得多。精度影响LoRA 的精度损失取决于 r 值秩r 值可训练参数微调精度损失适用场景817M0.5-1%简单任务分类1635M0.1-0.3%通用对话、代码3270M0.1%复杂任务翻译、推理64140M几乎无损失全量微调替代r16 是性价比最高的选择。除非任务特别复杂否则不需要 r32。LoRA 微调是昇腾NPU上最实用的微调方案——显存少、速度快、合并后零推理开销。如果你的场景是基座模型 领域适配LoRA 几乎总是比全量微调更好的选择。仓库在这里https://atomgit.com/cann/torch_npu