当前位置：首页 > news >正文

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

news 2026/5/23 20:19:45

本文基于昇腾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_LEVELS=np.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],dtype=np.float32)@staticmethoddefquantize(weight_fp16):""" weight_fp16: [out_dim, in_dim] FP16 权重 返回: uint8 数组（每个 uint8 装 2 个 4-bit 值） """# 对每个 1D 行做归一化——QLoRA 是逐行量化的shape=weight_fp16.shape w_flat=weight_fp16.flatten()# 算每行的 absmax——用来归一化到 [-1, 1]row_max=weight_fp16.abs().max(dim=-1,keepdim=True).values# [out_dim, 1]row_max=row_max.clamp(min=1e-12)# 归一化w_normalized=weight_fp16/row_max# 值范围 [-1, 1]# 映射到离最近的 NF4 levellevels=torch.tensor(NF4Quantizer.NF4_LEVELS,device=weight_fp16.device)indices=torch.bucketize(w_normalized,levels)-1indices=indices.clamp(0,15).to(torch.uint8)# 压缩：两个 4-bit 塞进一个 uint8packed=indices[...,::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] """levels=torch.tensor(NF4Quantizer.NF4_LEVELS)# 拆包lo=packed&0x0Fhi=(packed>>4)&0x0F# 交换使 shape 正确indices=torch.stack([lo,hi],dim=-1).reshape(shape)# 反量化：level[indices] * row_maxw_deq=levels[indices]*row_max.unsqueeze(-1)returnw_deq.to(torch.float16)

NF4 把 16-bit 权重压到 4-bit——省 4 倍显存。LLaMA-70B 从 140GB 压到 35GB，一张 Ascend 910（64GB）就能装下。

QLoRA 的前向流程

# QLoRA 的一层 Forward——冻结层反量化 + LoRA 分支 FP16classQLoRALayer(torch.nn.Module):def__init__(self,base_weight_fp16,lora_A,lora_B,nf4_packed,row_max,rank=8,alpha=16):super().__init__()# 冻结权重——以 NF4 格式存储，不参与梯度self.register_buffer("nf4_weight",nf4_packed)self.register_buffer("row_max",row_max)self.out_dim,self.in_dim=base_weight_fp16.shape# LoRA 适配器——FP16，参与训练self.lora_A=lora_A# [rank, in_dim]self.lora_B=lora_B# [out_dim, rank]self.scale=alpha/rank# 冻结的原始权重只在 Forward 时反量化# 不存反量化版本——省显存defforward(self,x):# Step 1: NF4 反量化——每次 Forward 都做# 实现里会用融合算子省掉搬来搬去w_deq=NF4Quantizer.dequantize(self.nf4_weight,self.row_max,(self.out_dim,self.in_dim))# Step 2: 原始路径——用反量化后的权重base_out=torch.nn.functional.linear(x,w_deq)# Step 3: LoRA 分支——保持 FP16 精度lora_out=self.lora_B(self.lora_A(x))*self.scalereturnbase_out+lora_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_ptr=gm_nf4+tile_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(inti=0;i<256;i+=2){uint8_tbyte=nf4_ptr[i/2];deq_values[i]=level_table[byte&0x0F];deq_values[i+1]=level_table[(byte>>4)&0x0F];}// 乘 row_max——恢复实际值范围for(intj=0;j<256;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 就能跑——显存占用约 42GB（35GB 量化权重 + 5GB LoRA + 2GB 中间 Tensor）。微调一个下游任务只需 6 小时，跟全参微调要 4 卡跑 3 天比，省了 40 倍资源。