""" ============================================================================= CS349E · 第 3 讲 · 02_outliers_llm_int8.py —— 离群特征为什么让 W8A8 崩溃，LLM.int8() 和 SmoothQuant 怎样救 ============================================================================= 【这个脚本要回答的问题】 W8A8：权重和激活都量化成 int8，用整数矩阵乘法计算（INT8 Tensor Core 的算力是 BF16 的 2 倍，字节也减半）。 LLM.int8()（Dettmers et al. 2022，本讲阅读材料）发现：模型到 6.7B 参数左右时，激活里会冒出"离群特征"—— 少数几个特征维度（约 0.1%）在几乎所有 token 上都很大（可达其他维度的 20 倍），8 位量化随之崩溃。 演示 1：在第 2 讲的模型里"注入"离群特征（不改变模型的输出），看激活变成什么样。 演示 2：复现论文表 1 的现象：没有离群特征时各种 int8 方案都还行；有离群特征时整张缩放、按 token 缩放、 vector-wise 都明显变差，只有 LLM.int8()（vector-wise + 混合精度分解）把输出救回来。 演示 3：为什么按 token / 按通道缩放管不住离群特征：缩放因子不能沿"被求和的维度"变化。 演示 4：阈值 α：多少维度要走 16 位，换来多少精度。 演示 5：SmoothQuant 的思路：把激活的"难度"挪给权重；以及它为什么对我们注入的离群特征完全免疫。 演示 6：LLM.int8() 的代价：在真实尺寸上，16 位的那部分占多少运算和字节。 【怎么读输出】偏差指标同 01 脚本。全部是确定性的计算，每次运行结果相同。 【运行】python3 02_outliers_llm_int8.py （约 3 秒） 【依赖】numpy；同目录的 quant_lab；../lecture01/console。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, print_header, print_observations, print_table # noqa: E402 from quant_lab import (QCFG, collect_act_absmax, compare, forward, init_params, inject_outliers, # noqa: E402 int8_symmetric, make_llm_int8_fn, make_smoothquant_fn, make_w8a8_fn, sample_tokens) OUTLIER_DIMS = [3, 77, 150] # 256 维里的 3 维（约 1.2%） OUTLIER_SCALE = 20.0 # 论文："可达其他维度的 20 倍" def setup(): params = init_params(QCFG, seed=0) toks = sample_tokens(params, QCFG, n_seq=8, length=64, seed=1) p_out = inject_outliers(params, OUTLIER_DIMS, OUTLIER_SCALE) return params, p_out, toks def fmt_cmp(r): return [f"{r['kl']:.5f}", f"{r['agree'] * 100:.1f}%", f"{r['ppl_ratio']:.4f}"] # --------------------------------------------------------------------------- # 演示 1：注入离群特征 # --------------------------------------------------------------------------- def demo1_inject(params, p_out, toks): print_header("演示 1：注入离群特征——模型的输出不变，激活的样子变了") d = np.max(np.abs(forward(params, QCFG, toks) - forward(p_out, QCFG, toks))) print(f"把第 {OUTLIER_DIMS} 维放大 {OUTLIER_SCALE:.0f} 倍（RMSNorm 之后），后面矩阵的对应行缩小同样倍数；" f"全精度 logits 的最大变化 {d:.1e}") rows = [] for label, pp in (("原模型", params), ("注入后", p_out)): st = collect_act_absmax(pp, QCFG, toks) for op in ("q_proj", "mlp_up", "o_proj", "mlp_down"): a = np.stack([st[(li, op)] for li in range(QCFG.n_layers)]) # [层, 特征维] n6 = int(np.max(np.sum(a >= 6.0, axis=1))) rows.append([label, op, f"{a.max():.1f}", f"{np.median(a):.2f}", f"{n6} / {a.shape[1]}"]) print_table(["模型", "线性层的输入", "最大 |x|", "各维度最大值的中位数", "最大值 ≥ 6 的维度数（最多的一层）"], rows) print_observations([ "注入是'换一种写法'：(x ⊙ γ)·W 里把 γ_j 放大、W 的第 j 行缩小，乘积不变，所以模型的输出一点没变（差 1e-14 量级）。" "但送进 Q/K/V 投影和 MLP 上投影的激活里，有 3 列变成了其他列的约 20 倍（最大 |x| 从 4.5 涨到 70）。", "这正是 LLM.int8() 论文描述的样子：离群值集中在**少数几个固定的特征维度**上，几乎每个 token 都有；" "出现在注意力投影和 MLP 上投影的输入里。（论文：6.7B 模型一个 2048 token 的序列里有约 15 万个离群值，" "却只集中在不超过 7 个特征维度上。）", "o_proj 和 mlp_down 的输入不受注入影响。注意 mlp_down 的输入（SwiGLU 的输出）本来尾巴就重：最大值 11.7，" "有不少维度超过 6——这是我们这个随机模型的特点，演示 4 会看到它的影响。", ]) # --------------------------------------------------------------------------- # 演示 2：复现论文表 1 的现象 # --------------------------------------------------------------------------- def demo2_table1(params, p_out, toks): print_header("演示 2：各种 int8 方案，有 / 没有离群特征（论文表 1 的现象）") methods = [ ("int8 整张 × 整张（absmax）", make_w8a8_fn(False, False)), ("int8 每 token × 整张", make_w8a8_fn(True, False)), ("int8 每 token × 每通道（vector-wise）", make_w8a8_fn(True, True)), ("LLM.int8()（vector-wise + 分解，α = 6）", make_llm_int8_fn(6.0)), ] rows = [] for name, fn in methods: cells = [name] for pp in (params, p_out): ref = forward(pp, QCFG, toks) r = compare(ref, forward(pp, QCFG, toks, linear_fn=fn), toks) cells += [f"{r['kl']:.5f}", f"{r['agree'] * 100:.1f}%"] rows.append(cells) print_table(["方案（lm_head 保持全精度）", "原模型 KL", "top-1", "有离群 KL", "top-1"], rows) print("\n论文表 1（C4 困惑度，越低越好）摘录：13B 模型 32 位 12.45；int8 absmax 19.08；vector-wise 16.48；" "LLM.int8() 12.45。6.7B：32 位 13.30；absmax 14.59；LLM.int8() 13.24。") print_observations([ "没有离群特征时，所有 int8 方案的 KL 都在 1e-4 量级：8 位对'正常'的激活绰绰有余。", "注入离群特征后，整张缩放的 KL 涨了约 28 倍（0.0008 → 0.023），按 token 缩放、vector-wise 也涨了十几倍——" "而模型本身一点没变。坏的不是模型，是**缩放因子被那 3 列'撑大'了**（演示 3）。", "LLM.int8() 把那几列挑出来用 16 位算，其余 99% 照样用 int8，KL 回到和没有离群特征时一样的水平。" "这和论文的结论一致：模型越大（离群特征出现后），只有 LLM.int8() 能保持全精度的困惑度。", ]) # --------------------------------------------------------------------------- # 演示 3：为什么按 token / 按通道缩放都管不住 # --------------------------------------------------------------------------- def demo3_why(params, p_out, toks): print_header("演示 3：为什么 vector-wise 管不住离群特征——缩放因子不能沿被求和的维度变化") rows = [] normal = np.setdiff1d(np.arange(QCFG.d_model), OUTLIER_DIMS) for label, pp in (("原模型", params), ("注入后", p_out)): cap = {} def grab(x, W, op, layer, cap=cap): if op == "q_proj" and layer == 0: cap["x"] = x.reshape(-1, x.shape[-1]).copy() return x @ W forward(pp, QCFG, toks, linear_fn=grab) X = cap["x"] q, s = int8_symmetric(X, per_row=True) err = (q * s - X)[:, normal] rel = np.linalg.norm(err) / np.linalg.norm(X[:, normal]) # 假想：每一列（特征维度）一个缩放因子 qc, sc = int8_symmetric(X.T, per_row=True) relc = np.linalg.norm((qc * sc).T[:, normal] - X[:, normal]) / np.linalg.norm(X[:, normal]) rows.append([label, f"{np.median(s) * 127:.2f}", f"{np.median(s):.4f}", f"{rel * 100:.2f}%", f"{relc * 100:.2f}%"]) print("第 0 层 q_proj 的输入 X（512 个 token × 256 维），只看 253 个'正常'维度的量化误差：") print_table(["模型", "每 token 最大 |x| 的中位数", "每 token 缩放因子（中位数）", "按 token 缩放：误差", "假想按列缩放：误差"], rows) print_observations([ "按 token 缩放时，每一行的最大值几乎总是落在离群那几列上（中位数从 3.0 涨到 24.8），缩放因子跟着变大约 8 倍：" "其余 253 个正常维度只能用很粗的格子量化，相对误差从 0.70% 涨到 6.4%。", "如果每一**列**（特征维度）一个缩放因子，离群列只影响自己，正常维度的误差回到 0.72%。可惜做不到：" "y_ij = Σ_k x_ik·w_kj，按列缩放意味着缩放因子随 k 变化，整数累加 Σ_k q_ik·q_kj 之后没法再统一乘回去。" "能放在矩阵乘法外面的，只有随 i（token）和 j（输出通道）变化的缩放因子——这就是 vector-wise 的极限。", "LLM.int8() 的思路：既然离群列只有几列，就把它们**整个拿出来**用 16 位算；SmoothQuant 的思路：在量化之前把" "每一列除以一个因子、再把它乘进权重的对应行（演示 5）。两者都绕开了'缩放因子不能沿 k 变化'的限制。", ]) # --------------------------------------------------------------------------- # 演示 4：阈值 α # --------------------------------------------------------------------------- def demo4_threshold(p_out, toks): print_header("演示 4：LLM.int8() 的阈值 α——多少维度走 16 位，换来多少精度（有离群特征的模型）") ref = forward(p_out, QCFG, toks) rows = [] for thr in (float("inf"), 12.0, 8.0, 6.0, 4.0, 3.0): log = [] r = compare(ref, forward(p_out, QCFG, toks, linear_fn=make_llm_int8_fn(thr, log)), toks) frac = sum(n for _, _, n, _ in log) / sum(d for _, _, _, d in log) inj = max(n for op, _, n, _ in log if op in ("q_proj", "mlp_up")) down = max(n for op, _, n, _ in log if op == "mlp_down") rows.append(["∞（纯 vector-wise）" if thr == float("inf") else f"{thr:g}", f"{frac * 100:.2f}%", inj, down] + fmt_cmp(r)) print_table(["α", "走 16 位的维度占比", "q/up 输入的离群维度", "down 输入", "KL 散度", "top-1 一致", "困惑度比"], rows) print_observations([ "α = 12 就已经恰好挑出注入的 3 个维度（q/up 输入），16 位部分只占全部维度的 0.65%，KL 就回到了 2.4e-4。", "α = 6（论文的取值）时，mlp_down 的输入也有约 10% 的维度被挑出来——不是因为它们是'系统性的离群特征'，" "而是这个随机模型的 SwiGLU 输出尾巴重。16 位部分涨到 3.7%，精度只多提高了一点点。真实大模型的正常激活要小得多，" "论文报告 α = 6 时只有约 0.1% 的维度走 16 位。", "α 再小，16 位的比例迅速上升（α = 3 时 67%），等于退回了 16 位矩阵乘法：精度好了，但 int8 的意义没了。" "阈值是一个'速度 vs 精度'的旋钮，要按模型实测来选。", ]) # --------------------------------------------------------------------------- # 演示 5：SmoothQuant # --------------------------------------------------------------------------- def demo5_smoothquant(params, p_out, toks): print_header("演示 5：SmoothQuant——把激活的难度挪给权重，再用最简单的整张 int8") rows = [] for label, pp in (("原模型", params), ("注入后", p_out)): ref = forward(pp, QCFG, toks) st = collect_act_absmax(pp, QCFG, toks) r0 = compare(ref, forward(pp, QCFG, toks, linear_fn=make_w8a8_fn(False, True)), toks) cells = [label, f"{r0['kl']:.5f}"] for a in (0.25, 0.5, 0.75, 1.0): r = compare(ref, forward(pp, QCFG, toks, linear_fn=make_smoothquant_fn(st, a)), toks) cells.append(f"{r['kl']:.5f}") rows.append(cells) print("激活整张一个缩放因子（静态，可以离线定好），权重每通道；表中是 KL 散度：") print_table(["模型", "不平滑", "α = 0.25", "α = 0.5", "α = 0.75", "α = 1.0"], rows) print_observations([ "SmoothQuant（Xiao et al. 2023，补充材料）：对每个特征维度 j 取 s_j = max|X_j|^α / max|W_j|^(1−α)，" "把 X 的第 j 列除以 s_j、W 的第 j 行乘以 s_j。(X/s)·(sW) = X·W 精确成立，但激活的离群列被压扁、难度挪给了权重。" "s 可以并进前一个 RMSNorm 的缩放系数里，运行时零开销；之后就能用最快的'整张一个缩放因子'的 int8 内核。", "注入离群特征后，不平滑的整张 int8 KL 是 0.0225（坏了）；平滑之后，每一列的数字都和原模型**完全相同**。" "原因是我们注入的方式恰好是'把 γ_j 放大、W_j 缩小'——SmoothQuant 只看 max|X_j|·max|W_j| 这个乘积，" "而注入不改变它。换句话说，SmoothQuant 正好能把这种离群特征'挪回去'。真实模型的离群特征不一定这么配合，" "论文在 OPT-175B 等模型上用 α = 0.5 左右达到了和 FP16 相当的 W8A8 精度。", "α 越大，挪给权重的越多：在这个模型里 α 从 0.25 到 1.0 一路变好（权重每通道量化，承受得住）。" "α 太大时权重会出现新的离群行，所以论文通常取 0.5 左右——又是一个要实测的旋钮。", ]) # --------------------------------------------------------------------------- # 演示 6：LLM.int8() 在真实尺寸上的代价 # --------------------------------------------------------------------------- def demo6_cost(): print_header("演示 6：LLM.int8() 的代价（真实尺寸，按论文'不超过 7 个离群维度'估算）") rows = [] for d, name in ((4096, "6.7B 级（d = 4096）"), (5120, "13B 级（d = 5120）"), (12288, "175B 级（d = 12288）")): n_out, n_o = 7, 4 * d # 以 MLP 上投影为例：W[d, 4d] frac = n_out / d w_int8 = d * n_o w_fp16_rows = n_out * n_o * 2 rows.append([name, n_out, f"{frac * 100:.2f}%", fmt_bytes(w_int8), fmt_bytes(w_fp16_rows), f"{w_fp16_rows / w_int8 * 100:.2f}%"]) print_table(["模型", "离群维度", "16 位部分的 FLOPs 占比", "int8 权重（一个 W[d, 4d]）", "离群行的 16 位拷贝", "额外字节"], rows) print_observations([ "16 位那部分的运算只占 0.06%–0.17%，离群行的 16 位拷贝只多 0.1%–0.3% 的字节：论文说'额外内存约 0.1%'，量级一致。", "真正的代价不在 FLOPs，而在**实现**：每次矩阵乘法都要先扫描激活找离群列、把矩阵拆成两块、跑两个内核、再把结果加起来。" "论文实测：6.7B 以下的模型用 LLM.int8() 反而比 FP16 慢；175B 尺寸的大矩阵乘法约快 2 倍。" "LLM.int8() 的主要目标是**省显存**（让 175B 模型装进一台消费级 GPU 服务器），而不是加速。", "【系统联系】'分解 + 两种精度'在硬件上不规整，这也是后来 SmoothQuant（离线平滑、单一 int8 内核）、" "FP8（动态范围大，离群值不那么致命）和块缩放格式（MX、NVFP4：每 16–32 个数一个缩放因子，离群值只影响自己那一块）" "更受欢迎的原因。", ]) if __name__ == "__main__": pa, po, tk = setup() demo1_inject(pa, po, tk) demo2_table1(pa, po, tk) demo3_why(pa, po, tk) demo4_threshold(po, tk) demo5_smoothquant(pa, po, tk) demo6_cost()