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 损失, 但在实践中面临三个核心挑战:
- 梯度稳定性: Gumbel-Softmax 噪声在深层模型上累积爆炸
- 显存瓶颈: 训练时中间激活占用远超推理
- 部署管线: 训练后如何快速打包并高效解包
本文提出 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-1B | 168 层, 679M 权重 |
| 量化 | DG 4-bit | group=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)
| 模式 | Fwd | GPU | Store | 技术 |
|---|---|---|---|---|
| bf16 | 35ms | 2.2GB | 2161MB | — |
| DG PyTorch unpack | 78ms | 1.3GB | 361MB | Python nibble |
| DG Triton unpack | 50ms | 1.2GB | 361MB | Triton kernel |
| DG materialized | 42ms | 5.1GB | — | bf16 缓存 |
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.3 | 29,000,000 | ❌ 噪声爆炸 |
| Gumbel-ST, TAU=3→0.5 | ~100,000 | ❌ 仍然失败 |
| Gumbel-ST, TAU=1→0.5 | 2,903 | ❌ 趋于崩溃 |
| STE | 60.0 | ✅ 稳定收敛 |
| STE, lr=5e-4 | 100.7 | 不稳定 |
| STE, lr=2e-4 | 60.0 | ✅ 最佳 |
| EN-only training | 48.1 | 最好 ppl |
| EN+CN mix | 60.0 | 中英平衡 |
4. 与现有方案对比
| DG 4-bit (Ours) | GPTQ | AWQ | GGUF/Q4_K | BitNet b1.58 | |
|---|---|---|---|---|---|
| 训练方式 | QAT (微调) | PTQ | PTQ | PTQ | 从头训练 |
| 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. 局限与未来工作
- 模型覆盖: 当前仅在 MiniCPM5-1B 上验证, 需扩展到 7B/13B 模型
- Water-filling 未启用: 为保 Triton 路径兼容性, 所有 block 使用统一 4-bit
- 生成速度: 50ms 是前向速度, 逐 token 生成时每 token 50ms = 20 tok/s
- 无 KV-cache 量化: 仅压缩权重, KV-cache 仍为 bf16
7. 结论
我们提出并验证了 DG 4-bit — 首个端到端 QAT 4-bit 量化管线。 核心贡献:
- Dyadic 码本+ block 分块架构: 理论简洁, 部署高效
- STE vs Gumbel-ST 系统性分析: 首次在 1B 模型 168 层上证明 STE 的必要性
- Triton 解包: 全模型 50ms, 1.2GB — 生产可用
- 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.