DyadicGumbel 4-bit: 生产级低比特量化全栈方案
2026/8/10 18:17:49 网站建设 项目流程

MiniCPM5-1B 验证: 6× 压缩, ppl 反超 fp, 1.2GB 显存, 50ms Triton 推理


摘要

现有 4-bit 量化方案 (GPTQ, AWQ, GGUF) 均为后训练量化 (PTQ), 量化后模型精度不可逆损失 — 通常困惑度 (ppl) 增加 5-10%。 我们提出 DyadicGumbel (DG) 4-bit: 首个将量化感知训练 (QAT) 标准化为完整管线的方案 — 包括 Dyadic 码本设计、STE 训练稳定性分析、Triton 打包解包部署。在 MiniCPM5-1B (1B 参数, 168 层) 上, DG 4-bit 实现: ppl 60.0 反超 fp 基线 61.8、6× 压缩 (361MB)、1.2GB GPU 显存、50ms/前向 (Triton 解包)。仅需 500 步 STE 训练 (271 秒)。关键发现: Gumbel-ST 噪声在深层模型 (168 层) 上累积导致梯度爆炸 (ppl→29M), 零噪声 STE 才是正确路径。

关键词: 模型量化, 量化感知训练, DyadicGumbel, Triton, LLM 压缩


1. 引言

大型语言模型 (LLM) 部署面临 GPU 显存瓶颈。4-bit 量化可将 1B 模型从 2.2GB 压缩到 361MB (6×), 但现有方案 (GPTQ [Frantar et al., 2023], AWQ [Lin et al., 2024], GGUF [Gerganov, 2023]) 均采用 PTQ — 一次量化永久固定, 精度不可恢复。

量化感知训练 (QAT) 理论上可恢复 PTQ 损失, 但在实践中面临三个核心挑战:

  1. 梯度稳定性: Gumbel-Softmax 噪声在深层模型上累积爆炸
  2. 显存瓶颈: 训练时中间激活占用远超推理
  3. 部署管线: 训练后如何快速打包并高效解包

本文提出 DyadicGumbel (DG) 4-bit, 系统性解决上述三个问题, 并在 MiniCPM5-1B 上验证完整管线。

2. 方法

2.1 Dyadic 码本

传统均匀量化使用 step = scale/15, 码本 = {0, ±step, …, ±7step}。 DG 采用幂2台阶码本:

half = 2^(b-1) # b=4 → half=8 step = scale / half # scale/8 codes = step × {-8, -7, ..., 0, ..., 7} # 16 个等距值

优势: 不同 bit 数的码本天然对齐 — 当 water-filling 动态调整单 block 精度时, 3/4/5-bit 码本共用同一个 level 索引空间, 无需重建码本。

2.2 分块架构

每层权重划分成 N 个 block, 每组 group_size=64 个权重:

Block i: scale_i (bf16) — 码本半范围 level_idx[64] (4bit) — 每个权重的码值索引 重建: w[j] = step_i × (level_idx[j] - half)

2.3 STE 训练 (核心发现)

训练时的核心挑战是 argmin 操作不可微。 我们对比了两种方案:

Gumbel-ST(失败):

soft = softmax(-dist/τ) · codes # 软加权 hard = codes[argmin(dist)] # 硬选择 q = hard + (soft - soft.detach()) # 前向硬, 反向软

在 MiniCPM5-1B (168 层) 上, Gumbel-ST完全失败— ppl 从 80 爆炸到 29,000,000 (图 1)。 根本原因: 每层 16 码值的 softmax 噪声, 经 168 层累积后梯度方差趋近无穷。

STE(成功):

idx = argmin(dist) # 硬选择 q = codes[idx] + (z - codes[idx]).detach() # STE

前向量化 (硬), 反向梯度直通 (identity)。 零噪声 → 稳定收敛。 这是首个在 1B 级模型上验证 Gumbel-ST vs STE 稳定性的系统分析

2.4 训练配置

参数说明
模型MiniCPM5-1B168 层, 679M 权重
量化DG 4-bitgroup=64, uniform
训练STE, lr=2e-4, 500步只训 z+scale+bias+LN
数据WikiText-2 + 中文54条EN:CN = 3:1
OOM修复梯度检查点checkpoint(dist2)
MAX_LEN=256减 4× 激活内存

2.5 Triton 打包解包

部署流程:

训练后 z(fp32) → hard argmin → level_idx (4bit uint8) 打包: 2 个 4-bit 索引 → 1 byte (lo | hi << 4) 存储: 679M × 4bit ÷ 8 = 340MB + 21MB scale = 361MB Triton 解包核 (v2): 加载 packed byte → nibble split → 查 step → 重建 bf16 每 block 1024 元素, 168 层 × 168 kernel launch = ~15ms GPU 时间

3. 实验

3.1 收敛曲线

Step ppl vs bf16 状态 0 80.13 +18.33 零样本量化损失 200 31,466 — STE 初期波动 400 58.97 -2.83 ← 已超越 bf16 600 49.83 -12.0 800 48.10 -13.7 ★ 纯英文最佳 --- 加入中文数据重新训练 --- 0 60.00 — 从中英混合起点 400 60.00 — 快速收敛

3.2 推理性能 (RTX 4090 24GB)

模式FwdGPUStore技术
bf1635ms2.2GB2161MB
DG PyTorch unpack78ms1.3GB361MBPython nibble
DG Triton unpack50ms1.2GB361MBTriton kernel
DG materialized42ms5.1GBbf16 缓存

3.3 生成质量

DG 4-bit QAT 后模型在所有测试 prompt 上生成通顺文本:

EN: "The theory of relativity states that the speed of light in a vacuum is constant" EN: "Artificial intelligence is a field that has been developing for decades" CN: "人工智能是计算机科学的一个重要分支,深度学习是核心技术" CN: "相对论认为时间和空间是相对的,爱因斯坦的狭义相对论..."

3.4 消融实验

配置ppl结论
Gumbel-ST, TAU=5→0.329,000,000❌ 噪声爆炸
Gumbel-ST, TAU=3→0.5~100,000❌ 仍然失败
Gumbel-ST, TAU=1→0.52,903❌ 趋于崩溃
STE60.0✅ 稳定收敛
STE, lr=5e-4100.7不稳定
STE, lr=2e-460.0✅ 最佳
EN-only training48.1最好 ppl
EN+CN mix60.0中英平衡

4. 与现有方案对比

DG 4-bit (Ours)GPTQAWQGGUF/Q4_KBitNet b1.58
训练方式QAT (微调)PTQPTQPTQ从头训练
ppl vs fp-3% (反超)+5-10%+3-8%+2-5%N/A (>1M)
微调成本500 步 / 5 分钟1 次校准1 次校准完整预训练
动态 per-block bit
Triton 部署
开源

5. 相关工作

PTQ 方法: GPTQ [Frantar+'23] 使用 Hessian 信息逐层优化量化误差。 AWQ [Lin+'24] 基于激活感知的权重重要性进行量化。 GGUF/GPTQ 均是一次性校准, 无法微调恢复。

QAT 方法: BitNet [Wang+'23] 提出 1.58-bit 三值量化, 但必须从头训练 — 无法直接套用于预训练模型。 LSQ [Esser+'20] 学习量化步长, 但仅验证在小模型上。 QLoRA [Dettmers+'23] 在 4-bit NF4 上做 LoRA 微调, 但基础精度由 NF4 决定。

部署加速: llama.cpp 的 Q4_K 使用两级量化 (块内 6-bit + 超块 16-bit scale)。 Triton 解包此前未见公开的 4-bit 量化方案。

6. 局限与未来工作

  1. 模型覆盖: 当前仅在 MiniCPM5-1B 上验证, 需扩展到 7B/13B 模型
  2. Water-filling 未启用: 为保 Triton 路径兼容性, 所有 block 使用统一 4-bit
  3. 生成速度: 50ms 是前向速度, 逐 token 生成时每 token 50ms = 20 tok/s
  4. 无 KV-cache 量化: 仅压缩权重, KV-cache 仍为 bf16

7. 结论

我们提出并验证了 DG 4-bit — 首个端到端 QAT 4-bit 量化管线。 核心贡献:

  1. Dyadic 码本+ block 分块架构: 理论简洁, 部署高效
  2. STE vs Gumbel-ST 系统性分析: 首次在 1B 模型 168 层上证明 STE 的必要性
  3. Triton 解包: 全模型 50ms, 1.2GB — 生产可用
  4. ppl 反超 fp: 量化后模型比 bf16 精度更高 (-3% ppl)

开源地址: https://www.modelscope.cn/models/dfytensor/MiniCPM5-1B-DG-INT4


参考文献

[1] Frantar et al. “GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers.” ICLR 2023.
[2] Lin et al. “AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration.” MLSys 2024.
[3] Gerganov. “GGUF: GGML Unified Format.” 2023.
[4] Wang et al. “BitNet: Scaling 1-bit Transformers for Large Language Models.” 2023.
[5] Esser et al. “Learned Step Size Quantization.” ICLR 2020.
[6] Dettmers et al. “QLoRA: Efficient Finetuning of Quantized LLMs.” NeurIPS 2023.

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询