量化感知训练(QAT)通过使用伪量化节点和直通估计器(STE)来模拟量化。这些技术被应用于将量化有效集成到对预训练大型语言模型(LLM)进行微调的实际操作流程中。标准微调使用模型原始的高精度权重(如FP32或BF16),使预训练模型适应特定的下游任务或数据集。QAT微调遵循类似原理,但有一个重要修改:模型学会执行任务,同时顾及低精度算术的影响。为QAT准备模型第一步是取一个预训练的、全精度模型,并为QAT做好准备。这需要修改模型架构,插入“伪量化”操作。如前所述,这些操作在正向传播时模拟量化和反量化值的效果,同时在反向传播时允许梯度不受影响地通过(或使用STE)。在此准备阶段,重要的决策包括:目标层: 模型中哪些层或组件应进行量化?通常,线性层(PyTorch中的nn.Linear)是主要选项,因为它们计算量大。嵌入层也可以考虑,尽管它们的量化有时更敏感。通常,靠近输入或输出的层,或通过性能分析显示敏感的层,可能会被排除在外或量化到更高精度。量化目标: 层的哪些部分进行量化?常见做法是量化权重(通常是静态的)和激活(通常在训练期间动态模拟,但最终使用习得的范围进行静态量化)。量化配置: 为每个目标层和组件(权重/激活)指定所需的位宽(例如INT8、INT4)、量化方案(对称或非对称)和粒度(逐张量、逐通道或逐组)。深度学习框架通常提供实用工具,根据配置文件或指定设置自动插入这些伪量化节点。例如,您可以指定所有线性层应使用对称逐通道量化将其权重量化为INT8,并使用非对称逐张量量化将其激活量化。QAT微调循环一旦模型准备好伪量化节点,微调过程就开始了。它与标准微调循环大致相同,但量化模拟处于活动状态。以下是典型训练步骤的分解:正向传播: 输入数据通过模型传播。当执行到达为QAT准备的层时:权重在用于操作(例如矩阵乘法)之前进行伪量化(先量化再反量化)。激活(操作或激活函数的输出)在传递到下一层之前进行伪量化。模型根据这些模拟的低精度操作计算输出和损失。反向传播: 损失梯度被计算并反向传播通过网络。直通估计器(STE)确保梯度可以流回伪量化节点,在梯度计算时有效忽略真实量化函数 $q(x)$ 的不可微特性。权重更新: 优化器根据计算出的梯度更新原始高精度权重。重要的一点是,模型学习到的权重值能够应对正向传播过程中模拟量化过程引入的噪声和信息损失。这一迭代过程使模型不仅能使其参数适应微调任务,还能适应目标低精度表示的限制。digraph QAT_Finetuning { rankdir=LR; node [shape=box, style=filled, fontname="Arial", color="#ced4da", fillcolor="#e9ecef"]; edge [fontname="Arial", color="#495057"]; subgraph cluster_layer { label = "QAT已启用层(例如,线性层)"; bgcolor="#f8f9fa"; style=dashed; node [shape=ellipse, fillcolor="#a5d8ff", color="#4dabf7"]; // 操作 inp [label="输入激活\n(FP32)", shape=box, fillcolor="#ffec99", color="#fcc419"]; wt_fp32 [label="原始权重\n(FP32)", shape=box, fillcolor="#ffec99", color="#fcc419"]; fake_quant_wt [label="伪量化\n(权重)", fillcolor="#fcc2d7", color="#f06595"]; fake_quant_act [label="伪量化\n(激活)", fillcolor="#fcc2d7", color="#f06595"]; op [label="线性操作\n(例如,矩阵乘法)", shape=cds, fillcolor="#b2f2bb", color="#51cf66"]; out [label="输出激活\n(FP32)", shape=box, fillcolor="#ffec99", color="#fcc419"]; wt_fp32 -> fake_quant_wt; fake_quant_wt -> op [label="模拟INT8权重"]; inp -> op; op -> fake_quant_act; fake_quant_act -> out [label="模拟INT8激活"]; // 反向传播提示 gradient [label="梯度流 (STE)", shape=plaintext, fontcolor="#7048e8"]; out -> gradient [style=dashed, arrowhead=none, color="#9775fa"]; gradient -> fake_quant_act [style=dashed, arrowhead=none, color="#9775fa"]; gradient -> op [style=dashed, arrowhead=none, color="#9775fa"]; gradient -> fake_quant_wt [style=dashed, arrowhead=none, color="#9775fa"]; gradient -> wt_fp32 [style=dashed, arrowhead=open, color="#9775fa"]; } inp_data [label="输入数据", shape=folder, fillcolor="#dee2e6"]; loss [label="计算损失", shape=invtriangle, fillcolor="#ffc9c9", color="#fa5252"]; optimizer [label="更新权重", shape=cylinder, fillcolor="#bac8ff", color="#5c7cfa"]; inp_data -> inp; out -> loss; loss -> gradient [style=invis]; // 反向流开始的逻辑连接 wt_fp32 -> optimizer [style=dashed, arrowhead=open, color="#5c7cfa"]; // 优化器更新权重 }简化了正向传播过程中QAT已启用层内的数据流,显示了伪量化节点的插入。反向传播使用STE更新原始全精度权重。训练考量使用QAT进行微调会带来一些特定考量:学习率: 由于模型正在适应任务和量化噪声,因此通常较低的学习率比标准微调更有利。这使得模型在调整适应模拟量化效果时能够更稳定地收敛。训练时长: QAT微调通常比从头开始训练模型所需的周期数更少,但可能需要比标准微调稍长的时间,以便模型权重能够稳定到适合量化的值。初始化: 从训练良好的FP32模型开始QAT微调是标准做法。有时,在引入伪量化节点之前,可能会进行几个周期的标准微调。批次大小: 批次大小的标准考量适用,以平衡计算资源和梯度稳定性。最终转换为整数模型需要记住,QAT微调过程的输出仍然是一个具有高精度权重和伪量化节点的模型。权重已经过优化,在模拟量化下表现良好,但模型本身尚未转换为适合部署的低精度整数格式。最后一步是将这个经过QAT微调的模型转换为一个真正的量化模型。此转换使用习得的量化参数(通常从QAT期间观察或习得的范围中导出)将优化的FP32权重转换为其INT8(或INT4等)表示。伪量化节点被实际的量化和反量化操作取代(如果硬件支持,则可融合为纯整数操作)。生成模型现在已准备好部署,充分利用了量化感知微调过程带来的精度优势。这一过程与PTQ(训练后通过单独的校准步骤确定量化参数,且权重未明确训练以抵抗量化噪声)形成对比。通过将这种感知集成到微调循环中,QAT为在最终量化模型中实现更高精度提供了途径,尤其是在较低位宽下。