""" ============================================================================= CS349E · 第 3 讲 · 01_weight_and_kv_quantization.py —— 只量化权重（和 KV cache）：比特数 vs 输出偏差 ============================================================================= 【这个脚本要回答的问题】 第 1、2 讲的结论：decode 的时间花在**搬字节**上——读权重、读 KV cache。最直接的办法是让每个数少占几个比特。 "仅权重量化"（weight-only，如 W4A16：权重 4 位、激活和计算仍是 16 位）不改变计算方式，只改变存储方式， 是最常用、最安全的推理量化。这里在第 2 讲的 Transformer 上量一量它的代价： 演示 1：各种格式（INT8 / FP8 / INT4 / MXFP4 / NVFP4，不同粒度）的"每个权重几比特"和输出偏差。 演示 2：INT4 的分组大小：组越小越准，但缩放因子的开销越大——收益递减。 演示 3：哪些矩阵最敏感：一次只量化一类矩阵（或一层），看谁贡献了最多的偏差。lm_head 是个例外。 演示 4：KV cache 量化：K、V 存成 FP8 / INT8 / INT4，输出偏多少；K 和 V 谁更敏感。 【怎么读输出】偏差的三个指标（定义见 quant_lab.py 开头）：KL 散度（越小越好）、top-1 一致率（越高越好）、 困惑度比（量化 / 全精度，越接近 1 越好）。"老师的文字"是全精度模型自己采样的 8 × 64 个 token。 全部是确定性的计算，每次运行结果相同。 【运行】python3 01_weight_and_kv_quantization.py （约 3 秒） 【依赖】numpy；同目录的 quant_lab（它再导入 ../lecture02/transformer_lab 和 ../prereq3/numfmt）；../lecture01/console。 """ import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import print_header, print_observations, print_table # noqa: E402 from quant_lab import (KV_SCHEMES, QCFG, WEIGHT_KEYS, WEIGHT_SCHEMES, QuantKVCache, compare, forward, # noqa: E402 init_params, kv_bits, quantize_int, quantize_params, quantize_weight, sample_tokens) def setup(): """全精度模型、老师采样的文字、全精度 logits（所有演示共用）。""" params = init_params(QCFG, seed=0) toks = sample_tokens(params, QCFG, n_seq=8, length=64, seed=1) return params, toks, forward(params, QCFG, toks) def fmt_cmp(r): """把 compare 的结果变成三个单元格。""" return [f"{r['kl']:.5f}", f"{r['agree'] * 100:.1f}%", f"{r['ppl_ratio']:.4f}"] # --------------------------------------------------------------------------- # 演示 1：各种格式 # --------------------------------------------------------------------------- def demo1_formats(params, toks, ref): print_header("演示 1：仅权重量化——各种格式的比特数与输出偏差（lm_head 保持全精度）") rows = [] for s in WEIGHT_SCHEMES: pq, bits = quantize_params(params, s) rows.append([s, f"{bits:.3f}"] + fmt_cmp(compare(ref, forward(pq, QCFG, toks), toks))) print(f"模型 {QCFG.name}：{QCFG.n_layers} 层，d = {QCFG.d_model}，d_ff = {QCFG.d_ff}；全精度困惑度 " f"{compare(ref, ref, toks)['ppl_ref']:.2f}") print_table(["格式", "比特 / 权重", "KL 散度", "top-1 一致", "困惑度比"], rows) print_observations([ "8 位几乎无损：INT8（每通道）的 KL 只有 7e-5、top-1 一致 99%。权重的 8 位量化是'免费的午餐'——字节减半，输出几乎不变。", "同样 8 位，FP8（E4M3，整张一个缩放因子）比 INT8 差：E4M3 只有 3 位尾数，大约 1/16 的相对精度；" "INT8 对'大致均匀分布在一个范围里'的权重更细。FP8 的优势在于动态范围大（对激活更友好），不在于精度。", "降到 4 位，偏差上了两个数量级，**粒度变得至关重要**：整张一个缩放因子（int4-tensor）KL 0.053，" "每通道 0.024，每 32 个一组 0.014。块缩放格式也遵循同样的规律：MXFP4（4.25 位）≈ INT4 每 128 个一组，" "NVFP4（4.5 位）≈ INT4 每 32 个一组。", "困惑度比和 KL 大体同向，但只用了 504 个 token，有噪声（例如 MXFP4 的困惑度比比 int4-g128 高，KL 却差不多）。" "真实评测要用大得多的数据集。", ]) # --------------------------------------------------------------------------- # 演示 2：INT4 的分组大小 # --------------------------------------------------------------------------- def demo2_group_size(params, toks, ref): print_header("演示 2：INT4 分组大小——更细的缩放因子 vs 更多的开销") rows = [] for g in (16, 32, 64, 128, 256): new = {k: v for k, v in params.items() if k != "layers"} new["layers"] = [] for lp in params["layers"]: q = dict(lp) for k in WEIGHT_KEYS: q[k] = quantize_int(lp[k].T, 4, granularity="group", group=g)[0].T new["layers"].append(q) bits = 4 + 16.0 / g rows.append([g, f"{bits:.4f}", f"{(bits - 4) / bits * 100:.1f}%"] + fmt_cmp(compare(ref, forward(new, QCFG, toks), toks))) print_table(["组大小", "比特 / 权重", "缩放因子占比", "KL 散度", "top-1 一致", "困惑度比"], rows) print_observations([ "组越小，KL 越小（16 个一组 0.011，256 个一组 0.023），但每个缩放因子（16 位）分摊到的权重越少：16 个一组时 20% 的比特花在缩放因子上。", "收益递减：组大小从 256 减到 128 再到 64，KL 每次只降 15%–17%，比特却在涨。业界常用 128（GPTQ、AWQ 的默认值）" "或块缩放格式的 16 / 32——这是一个'精度 vs 字节'的工程折中，没有唯一正确的答案。", "这里只用了最简单的'就近舍入'（RTN）。GPTQ 用二阶信息逐列修正舍入误差、AWQ 按激活大小保护重要通道，" "同样的比特能做得更准——它们改变的是'怎么舍入'，不改变'存多少字节'。", ]) # --------------------------------------------------------------------------- # 演示 3：哪些矩阵最敏感 # --------------------------------------------------------------------------- def demo3_sensitivity(params, toks, ref): print_header("演示 3：一次只量化一部分（INT4，每 128 个一组）——谁贡献了偏差") rows = [] names = {"wq": "Wq", "wk": "Wk", "wv": "Wv", "wo": "Wo", "w_gate": "W_gate", "w_up": "W_up", "w_down": "W_down"} total_w = sum(lp[k].size for lp in params["layers"] for k in WEIGHT_KEYS) for k in WEIGHT_KEYS: new = {kk: v for kk, v in params.items() if kk != "layers"} new["layers"] = [] for lp in params["layers"]: q = dict(lp) q[k] = quantize_weight(lp[k], "int4-g128")[0] new["layers"].append(q) share = sum(lp[k].size for lp in params["layers"]) / total_w rows.append([f"只量化每层的 {names[k]}", f"{share * 100:.1f}%"] + fmt_cmp(compare(ref, forward(new, QCFG, toks), toks))) new = dict(params) new["lm_head"] = quantize_weight(params["lm_head"], "int4-g128")[0] rows.append(["只量化 lm_head", f"({params['lm_head'].size / total_w * 100:.1f}%)"] + fmt_cmp(compare(ref, forward(new, QCFG, toks), toks))) for li in range(QCFG.n_layers): new = {kk: v for kk, v in params.items() if kk != "layers"} new["layers"] = [dict(lp) for lp in params["layers"]] for k in WEIGHT_KEYS: new["layers"][li][k] = quantize_weight(params["layers"][li][k], "int4-g128")[0] rows.append([f"只量化第 {li} 层的全部矩阵", f"{100 / QCFG.n_layers:.1f}%"] + fmt_cmp(compare(ref, forward(new, QCFG, toks), toks))) pq, _ = quantize_params(params, "int4-g128") rows.append(["全部线性层（不含 lm_head）", "100%"] + fmt_cmp(compare(ref, forward(pq, QCFG, toks), toks))) print_table(["量化了什么", "占线性层权重", "KL 散度", "top-1 一致", "困惑度比"], rows) print_observations([ "MLP 的三个矩阵各自贡献的偏差比注意力的四个大（它们也更大：每个是 Wq 的 3 倍）；" "Wq、Wk 最不敏感——它们的误差先经过 softmax 才影响输出。", "单独量化 lm_head（它不在'线性层'里，参数量相当于线性层的 9%）造成的 KL（0.027）比量化**其余全部**线性层（0.020）还大：" "lm_head 的误差直接加在 logits 上，没有任何后续层来'消化'。实践中常把 lm_head（和词嵌入）留在 8 / 16 位。", "四层的敏感度差不多，越靠前略高（第 0 层 0.0058，第 3 层 0.0043）。各部分的 KL 加起来约等于全部量化时的 KL：" "小扰动下，误差近似可加。这正是'混合精度'策略的基础：按敏感度给不同部分分配不同的比特。", ]) # --------------------------------------------------------------------------- # 演示 4：KV cache 量化 # --------------------------------------------------------------------------- def demo4_kv(params, toks, ref): print_header("演示 4：KV cache 量化（每个 token、每个 KV 头一个缩放因子）") rows = [] for ks, vs in [(s, s) for s in KV_SCHEMES] + [("int4", "bf16"), ("bf16", "int4")]: cache = QuantKVCache(QCFG, toks.shape[0], ks, vs) r = compare(ref, forward(params, QCFG, toks, cache=cache), toks) bits = (kv_bits(ks, QCFG.d_head) + kv_bits(vs, QCFG.d_head)) / 2 rows.append([f"K {ks} / V {vs}", f"{bits:.1f}", f"{16 / bits:.2f}×"] + fmt_cmp(r)) print_table(["KV 格式", "比特 / 数", "KV 缩小", "KL 散度", "top-1 一致", "困惑度比"], rows) print_observations([ "INT8 的 KV 几乎无损（KL 1e-5）；FP8 稍差（KL 1e-4，原因同演示 1：尾数只有 3 位）。KV 字节减半，" "长上下文、大批量时 decode 的时间几乎减半（KV 是那时的主要字节，第 2 讲 §5.6），能放下的序列数也翻倍。", "INT4 的 KV 偏差大一个数量级（KL 0.002）。在这个随机模型里 V 比 K 更敏感（只量化 V 的 KL 是只量化 K 的 2 倍）；" "真实模型里常常相反——K 有离群通道（类似演示 2 的激活离群），所以 KIVI 等方法让 K 按通道、V 按 token 分组。" "**哪个更敏感取决于模型，要测，不要猜**。", "【系统联系】KV 量化和权重量化省的是不同的字节：小批量时权重是大头，大批量长上下文时 KV 是大头（03 脚本）。", ]) if __name__ == "__main__": p, t, r = setup() demo1_formats(p, t, r) demo2_group_size(p, t, r) demo3_sensitivity(p, t, r) demo4_kv(p, t, r)