QLoRA：4-bit 量化微调的完整链路

本文基于昇腾CANN和昇腾NPU围绕 cann-recipes-train 仓库的相关技术展开。QLoRA 不是简单的 LoRA 量化。它在 LoRA 的冻结权重上做了 NF4 量化同时保留了 LoRA 适配器的 FP16 精度。CANN 上部署 QLoRA 模型时NF4 的反量化要在 NPU 上做不能让 CPU 插一手。NF4 量化怎么把权重压到 4-bit# NF4 量化——正态分布的 4-bit 量化importtorchimportnumpyasnpclassNF4Quantizer: NF4: Normal Float 4——值分布按正态分布的百分位分桶 16 个桶每个桶有相同概率正态下 所以值密集的地方桶多稀疏的地方桶少 # NF4 的 16 个量化值——从标准正态分布 CDF 的等间隔百分位算出NF4_LEVELSnp.array([-1.0,-0.6962,-0.5251,-0.3926,-0.2779,-0.1728,-0.0739,0.0000,0.0739,0.1728,0.2779,0.3926,0.5251,0.6962,1.0000,1.5000],dtypenp.float32)staticmethoddefquantize(weight_fp16): weight_fp16: [out_dim, in_dim] FP16 权重 返回: uint8 数组每个 uint8 装 2 个 4-bit 值 # 对每个 1D 行做归一化——QLoRA 是逐行量化的shapeweight_fp16.shape w_flatweight_fp16.flatten()# 算每行的 absmax——用来归一化到 [-1, 1]row_maxweight_fp16.abs().max(dim-1,keepdimTrue).values# [out_dim, 1]row_maxrow_max.clamp(min1e-12)# 归一化w_normalizedweight_fp16/row_max# 值范围 [-1, 1]# 映射到离最近的 NF4 levellevelstorch.tensor(NF4Quantizer.NF4_LEVELS,deviceweight_fp16.device)indicestorch.bucketize(w_normalized,levels)-1indicesindices.clamp(0,15).to(torch.uint8)# 压缩两个 4-bit 塞进一个 uint8packedindices[...,::2]|(indices[...,1::2]4)returnpacked.cpu().numpy(),row_max.cpu().numpy()staticmethoddefdequantize(packed,row_max,shape): packed: 量化后的 uint8 数组 row_max: 每行的 absmax shape: 原始 [out_dim, in_dim] levelstorch.tensor(NF4Quantizer.NF4_LEVELS)# 拆包lopacked0x0Fhi(packed4)0x0F# 交换使 shape 正确indicestorch.stack([lo,hi],dim-1).reshape(shape)# 反量化level[indices] * row_maxw_deqlevels[indices]*row_max.unsqueeze(-1)returnw_deq.to(torch.float16)NF4 把 16-bit 权重压到 4-bit——省 4 倍显存。LLaMA-70B 从 140GB 压到 35GB一张 Ascend 91064GB就能装下。QLoRA 的前向流程# QLoRA 的一层 Forward——冻结层反量化 LoRA 分支 FP16classQLoRALayer(torch.nn.Module):def__init__(self,base_weight_fp16,lora_A,lora_B,nf4_packed,row_max,rank8,alpha16):super().__init__()# 冻结权重——以 NF4 格式存储不参与梯度self.register_buffer(nf4_weight,nf4_packed)self.register_buffer(row_max,row_max)self.out_dim,self.in_dimbase_weight_fp16.shape# LoRA 适配器——FP16参与训练self.lora_Alora_A# [rank, in_dim]self.lora_Blora_B# [out_dim, rank]self.scalealpha/rank# 冻结的原始权重只在 Forward 时反量化# 不存反量化版本——省显存defforward(self,x):# Step 1: NF4 反量化——每次 Forward 都做# 实现里会用融合算子省掉搬来搬去w_deqNF4Quantizer.dequantize(self.nf4_weight,self.row_max,(self.out_dim,self.in_dim))# Step 2: 原始路径——用反量化后的权重base_outtorch.nn.functional.linear(x,w_deq)# Step 3: LoRA 分支——保持 FP16 精度lora_outself.lora_B(self.lora_A(x))*self.scalereturnbase_outlora_out# Forward 做了 1 次反量化 1 次 FP16 MatMul 2 次小 MatMul# 反量化的开销约 0.03ms——比读显存省的时间划算CANN 上的 NF4 融合算子// Ascend C 实现的 NF4 反量化 MatMul 融合——省掉反量化写回classNF4MatMulKernel:publicAscendC::Kernel{__aicore__inlinevoidProcess()override{// Step 1: 加载量化权重——4-bit每次 Tile 读 256 个 NF4 值// 256 个 NF4 值 128 bytes比 FP16 版本的 512 bytes 小 4 倍uint8_t*nf4_ptrgm_nf4tile_offset;// Step 2: 在 L1 上做反量化// 按 level 表查表——用 L1 的 Lookup Table 指令floatlevel_table[16]{-1.0,-0.6962,...,1.5};// 拆包两个 4-bit 取出来// 查表生成 FP16 值——直接在 Vector Unit 上做float16_t deq_values[256];for(inti0;i256;i2){uint8_tbytenf4_ptr[i/2];deq_values[i]level_table[byte0x0F];deq_values[i1]level_table[(byte4)0x0F];}// 乘 row_max——恢复实际值范围for(intj0;j256;j){deq_values[j]*row_max_val;}// Step 3: 反量化完的数据直接进 Cube——不写回 DDR// 省掉 dequantize → DDR → MatMul 的两趟搬运AscendC::MatMul(output,input_local,deq_values,AscendC::CUBE_MATRIX_TYPE::NORMAL);}};QLoRA 在显存受限场景下特别值。LLaMA-70B 用 QLoRA 微调时单卡 Ascend 910 就能跑——显存占用约 42GB35GB 量化权重 5GB LoRA 2GB 中间 Tensor。微调一个下游任务只需 6 小时跟全参微调要 4 卡跑 3 天比省了 40 倍资源。参考仓库QLoRA 微调示例TorchAir 量化微调支持pyasc 量化工具