十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Needle 2的QAT量化感知训练:fake_quant、STE与CQ噪声注入时机全解

Needle 2的QAT量化感知训练:fake_quant、STE与CQ噪声注入时机全解 Needle 2的QAT量化感知训练fake_quant、STE与CQ噪声注入时机全解【免费下载链接】needle14MB foundation model for tiny devices; phones, wearables, smart home, and robots.项目地址: https://gitcode.com/GitHub_Trending/needle20/needleNeedle 2 是一款面向手机、可穿戴设备等微型终端的 14MB 基础模型而 QAT量化感知训练正是它能把 45M 参数压缩到 2-bit、全会话仅占 28MB 内存的核心技术。本文结合 needle/model/quantize.py 的真实实现用零门槛的方式讲清 QAT 的三大主角——fake_quant仿真量化、STE 直通估计器、CQ 噪声注入——分别在什么时机生效、为什么这样设计帮助新手快速看懂这套边训练、边模拟量化的完整流程。为什么 Needle 2 必须靠 QAT从 45M 参数到 14MB 二进制先说结论QAT 解决的是低比特模型不聪明这个老大难问题。直接把 fp16 模型四舍五入到 2-bit这叫训练后量化PTQ精度往往崩塌。QAT 的思路则是在训练阶段就模拟量化误差让模型在带噪的世界里学会生存导出时再真正量化——精度损失被降到最低。上图正是 QAT 的价值证明Needle 2 以CQ2-bitCactus Quants 2 比特量化后在 Mobile-Actions 基准上打平 FunctionGemma 270M、LFM2.5 230M 这类 570 倍大的 fp16 模型——这不是靠硬件而是靠 QAT 把量化损失训没了。QAT 三件套一张表看懂分工在动手之前先看全景。needle/model/quantize.py 中三个关键机制各司其职机制角色作用对象生效时机fake_quant前向仿真量化按 128 组 absmax 量化到 4-bit权重 kernel / embedding训练前向传播时按QAT_EVERY频率触发STE直通估计器让梯度绕过不可导的 round 操作所有仿真正/量化操作反向传播时CQ 噪声注入按 CQ 失真率标定强度向权重加高斯噪声可量化张量训练时替代/配合硬量化fake_quant用彩排提前体验量化核心函数在 needle/model/quantize.py#L10-L22它做四件事分组把权重按group_size128切成小组维度不够就零填充再切定标每组取absmax求缩放因子scale absmax / qmaxqmax 2^(bits-1) - 14-bit 时为 7量化-反量化clip(round(w / scale)) * scale——这一来一回连续权重被掰到最近的量化格点STE 缝合最后返回w stop_gradient(q - w)。第 4 步就是整个 QAT 最巧的一笔前向传播看到量化后的值反向传播看到的却是原权重——误差项(q - w)被stop_gradient冻结梯度像没发生量化一样直接穿过。这就是所谓的STEStraight-Through Estimator直通估计器。为什么需要 STEround()几乎处处不可导如果梯度在这里断掉量化前的权重永远收不到更新信号。STE 用一个梯度恒等的替身把路铺通模型在带量化误差的景观里前向却沿着平滑的原始路径回传梯度——两全其美。同样的套路在部署侧也有镜像cq_steneedle/model/quantize.py#L354-L355用w stop_gradient(cq_quantize(w) - w)包住 CQ 码本量化保证评估/混合精度路径同样可导。CQ 噪声注入注入什么、多少、何时fake_quant模拟的是格点量化而 Needle 2 部署时用的是另一套CQCactus QuantsWalsh-Hadamard 旋转变换 Lloyd-Max 高斯码本见cq_quantizeneedle/model/quantize.py#L134-L147甚至支持 1.58-bit 三值化。两种量化的失真特性不同所以噪声注入必须按部署方案定剂量剂量怎么算noise_scale(bits)先用cq_distortion在随机矩阵上实测 CQ 的相对失真率再对比特数做对数插值最终sigma 组内 RMS × scale。也就是说注入的噪声强度精确等于部署后真实会有的量化误差量级怎么加add_cq_noiseneedle/model/quantize.py#L242-L251对每个 128 组独立采样高斯噪声noise_params则批量作用到所有可量化张量kernel、embedding、mhc_phi。注入时机是 QAT 与普通训练的分水岭训练期用噪声或周期性fake_quant替代硬量化。权重每次更新都在抖动模型被迫学会对量化级误差鲁棒——这正是量化感知四个字的含义部署期噪声全部撤掉deploy_quantize直接执行真实的 CQ 量化产出.cact归档。权重量化本身的触发频率由configure_qat(every, ...)控制needle/model/quantize.py#L56-L58maybe_quant_weights通过jax.lax.cond按QAT_EVERY决定这一步是否真的过一遍 4-bit 组量化组大小 128。这种间歇彩排比每步都量化省算力又足够让模型适应量化景观。激活 8-bit 与 KV 缓存量化推理路径上的注入点除了权重前向路径上的激活也被纳入 QAT。在 needle/model/architecture.py 中_aq()L22-L25用jax.lax.cond(quant, fake_quant_act, ...)把 8-bit 激活量化ACT_BITS 8精确地挂在四个位置embedding 读出后每个注意力块内部残差、输出投影最终 logits 之前MTP多 token 预测分支的拼接处。KV 缓存同样有一处量化闸门maybe_quant_kvneedle/model/quantize.py#L41-L45在KV_BITS配置非零时对 K/V 做 CQ 64 组 fake 量化——这是 256-token 滑窗推理时内存控制的关键。而configure_deploy会在比特配置变化时自动jax.clear_caches()避免 JIT 缓存串味。如上图所示QAT 的量化点正好覆盖每个 Transformer 块的激活主干——权重 4-bit组量化、激活 8-bit、KV 可选 8-bit三者共同构成部署时的完整量化配方也即导出信息里常见的W4A8 / CQ W4A8标记。从训练到部署两套仿真如何闭环整个生命周期的衔接在 needle/model/finetune.py 的build_main里完成微调needle finetune用 LoRA 微调冻结的基座默认 rank 16 / alpha 32AdamW warmup-cosine 调度细节见 doc/finetuning.md构建needle build合并 LoRA 适配器再按检查点声明的逐层比特映射mixed bits如敏感层 3-bit、其余 2-bit走cq_quantize_params真实量化输出单个.cact部署14MB 引擎直接加载.cact训练时的fake_quant/噪声仿真与部署时的 CQ 量化由同一套cq_quantize实现保证仿真误差 真实误差。这也是 QAT 与 PTQ 的本质区别PTQ 训完才量化误差无处可逃QAT 训时就量化或等价噪声误差被模型内化。新手快速上手三步跑通 QAT 微调流程pip install cactus-needle needle finetune data.jsonl --epochs 3 needle build checkpoints/needle2.pkl --lora checkpoints/needle_lora.pkl --out my.cact数据格式JSONL一行一个{query, tools, answers}示例doc/finetuning.md 有完整说明--bits 2|4或检查点内嵌比特映射控制量化宽度--upload可发布归档运行时用needle.Needle(weightsmy.cact, tools[...])加载引擎对权重无感知推理行为契约与 API 细节见 doc/apis.md权重加载机制的回归测试可参考 tests/test_weights.py。总结三个关键时机一句话记住前向时——fake_quant按 128 组 4-bit absmax 把权重掰到量化格点激活走 8-bit反向时——STE 用stop_gradient缝合梯度无视量化墙直接穿过训练全程——CQ 噪声按部署失真率标定注入让模型在噪声中学出鲁棒性部署时撤掉噪声、执行真量化零落差上线。这三件事缺一不可没有 fake_quant模型没体验过量化没有 STE梯度断流训不动没有 CQ 噪声仿真与部署之间就隔着一条精度鸿沟。读懂 needle/model/quantize.py 这不到 400 行代码你就掌握了微型端侧模型 QAT 的完整方法论。【免费下载链接】needle14MB foundation model for tiny devices; phones, wearables, smart home, and robots.项目地址: https://gitcode.com/GitHub_Trending/needle20/needle创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表