""" ============================================================================= CS349E · 第 3 讲 · 04_sparsity.py —— 稀疏：非结构化、2:4 结构化、激活稀疏，各省了什么 ============================================================================= 【这个脚本要回答的问题】 量化让每个数变小；稀疏（sparsity）让一部分数**直接变成 0、不存不算**。听起来更彻底，但有两个难点： 剪掉哪些（精度），以及怎样存、怎样算才能真的省下时间（硬件）。 演示 1：剪掉 50% 的权重：非结构化（哪里小剪哪里）vs 2:4 结构化（每 4 个里留 2 个），以及和 INT8 / INT4 量化的比较。 演示 2：存储格式：非结构化稀疏要存"非零元素在哪"，稀疏度不够高时反而比稠密还大；2:4 的元数据只要每个非零 2 比特。 演示 3：2:4 在 roofline 上：Tensor Core 的稀疏模式让算力翻倍（帮 prefill），字节省到 9/16（帮 decode）——和量化比一比。 演示 4：激活稀疏：ReLU 让 MLP 的一部分神经元输出 0，batch = 1 时可以少读 W_down 的对应行；可批一大，省下的就没了。 【怎么读输出】偏差指标同 01 脚本（KL、top-1 一致、困惑度比）。全部是确定性的计算，每次运行结果相同。 【运行】python3 04_sparsity.py （约 2 秒） 【依赖】numpy；同目录的 quant_lab；../lecture01 的 console、hw_specs、llm_math。 """ 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_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 from llm_math import LLAMA31_8B # noqa: E402 from quant_lab import (QCFG, WEIGHT_KEYS, compare, forward, init_params, quantize_params, sample_tokens, # noqa: E402 step_time) def prune_unstructured(W, frac): """ 非结构化剪枝：整个矩阵里绝对值最小的 frac 比例的权重置 0。 【返回】剪枝后的矩阵 【边界】并列时（多个数一样小）np.partition 会任意挑，影响可以忽略。 【Python 小知识】np.partition(a, k) 只保证第 k 小的数放在第 k 位、左边都更小，比完整排序快。 """ k = int(W.size * frac) if k == 0: return W.copy() thr = np.partition(np.abs(W).ravel(), k - 1)[k - 1] return np.where(np.abs(W) > thr, W, 0.0) def prune_2_4(W): """ 2:4 结构化剪枝：沿**输入维**（点积的方向）每连续 4 个权重一组，只保留绝对值最大的 2 个。 【为什么沿输入维】NVIDIA Ampere 以后的稀疏 Tensor Core 要求矩阵乘法里被求和的那一维满足"每 4 个最多 2 个非零"， 硬件按 2 比特的下标去取对应的激活，跳过另外一半的乘法。 【参数】W：[输入维, 输出维]，输入维必须是 4 的倍数 【返回】剪枝后的矩阵 【Python 小知识】np.argsort(..., axis=-1) 给出每组内的排序下标；np.put_along_axis 按下标把值写回去。 """ n_in, n_out = W.shape g = W.T.reshape(n_out, n_in // 4, 4) order = np.argsort(np.abs(g), axis=-1) # 每组 4 个从小到大 mask = np.ones_like(g, dtype=bool) np.put_along_axis(mask, order[..., :2], False, axis=-1) return (g * mask).reshape(n_out, n_in).T def prune_params(params, method): """对所有线性层（不含 lm_head）剪枝；method 为 "unstructured" 或 "2:4"。""" 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] = prune_unstructured(lp[k], 0.5) if method == "unstructured" else prune_2_4(lp[k]) new["layers"].append(q) return new def setup(): 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) # --------------------------------------------------------------------------- # 演示 1：剪掉 50% # --------------------------------------------------------------------------- def demo1_prune(params, toks, ref): print_header("演示 1：剪掉 50% 的权重——非结构化 vs 2:4，以及和量化比") W = np.concatenate([lp[k].ravel() for lp in params["layers"] for k in WEIGHT_KEYS]) rows = [] for label, pp, bits in (("非结构化 50%（按幅度）", prune_params(params, "unstructured"), "16 + 下标"), ("2:4 结构化（按幅度）", prune_params(params, "2:4"), "9.0"), ("INT8 每通道（对照）", quantize_params(params, "int8-channel")[0], "8.05"), ("INT4 每 128 个一组（对照）", quantize_params(params, "int4-g128")[0], "4.13")): Wp = np.concatenate([lp[k].ravel() for lp in pp["layers"] for k in WEIGHT_KEYS]) rel = np.linalg.norm(Wp - W) / np.linalg.norm(W) r = compare(ref, forward(pp, QCFG, toks), toks) rows.append([label, bits, f"{rel * 100:.1f}%", f"{r['kl']:.5f}", f"{r['agree'] * 100:.1f}%", f"{r['ppl_ratio']:.4f}"]) print_table(["做法（所有线性层，lm_head 不动）", "比特 / 权重", "权重相对误差", "KL 散度", "top-1 一致", "困惑度比"], rows) print_observations([ "剪掉一半权重，非结构化的相对误差是 26.7%：被剪掉的是最小的那一半，但它们加起来仍占了约 7% 的'能量'（26.7% 的平方）。" "2:4 的限制更多（每 4 个里必须剪 2 个，哪怕 4 个都很大），误差明显更大（36.4%）。", "在这个随机模型上，50% 稀疏的输出偏差比 INT4 量化大得多（KL 0.09–0.16 vs 0.02），比特数却是 INT4 的两倍多。" "随机高斯权重没有'可以安全剪掉'的冗余；真实训练过的模型有，配合 SparseGPT、Wanda 这类按激活加权的剪枝方法" "和少量微调，2:4 的精度损失可以小得多——但通常仍比 8 位量化大。", "**稀疏和量化不是二选一**：2:4 的非零值本身也可以是 FP8 / INT8（硬件支持 2:4 + FP8），两种压缩可以叠加。", ]) # --------------------------------------------------------------------------- # 演示 2：存储格式 # --------------------------------------------------------------------------- def demo2_storage(): print_header("演示 2：存稀疏矩阵要多少比特——'非零在哪'本身也要存") rows = [] for s in (0.5, 0.75, 0.9, 0.99): nnz = 1 - s csr = nnz * (16 + 16) # 每个非零：16 位的值 + 16 位的列下标（行指针分摊后可忽略） bitmap = nnz * 16 + 1 # 每个非零 16 位的值 + 每个位置 1 位的掩码 rows.append([f"{s * 100:.0f}%", "16", f"{csr:.2f}", f"{bitmap:.2f}", "9.00" if s == 0.5 else "-"]) print_table(["稀疏度", "稠密 FP16", "CSR（值 + 16 位列下标）", "位图（值 + 每位置 1 位）", "2:4（值 + 每个非零 2 位）"], rows) print_observations([ "50% 稀疏时，CSR 格式每个权重平均要 16 位——和稠密一样大，一个字节都没省：省下的值，全被'下标'吃掉了。" "非结构化稀疏要到 75%–90% 以上才开始在存储上划算。", "2:4 的下标只要 2 位（在 4 个位置里选 2 个），每个权重平均 (2×16 + 2×2) / 4 = 9 位，是稠密的 56%。" "规整的结构不仅让硬件好算，也让元数据便宜。", "【系统联系】这和第 1 讲的强度账一脉相承：稀疏省下的是'运算'，但如果为了跳过运算要多读下标、做不规则的访存，" "在 memory-bound 的 decode 里可能得不偿失。**省 FLOPs 不等于省时间**。", ]) # --------------------------------------------------------------------------- # 演示 3：2:4 在 roofline 上 # --------------------------------------------------------------------------- def demo3_roofline(): print_header("演示 3：2:4 稀疏在 roofline 上——Llama-3.1-8B，H100（稀疏 Tensor Core 的峰值是稠密的 2 倍）") m, chip = LLAMA31_8B, H100 rows = [] cases = [("稠密 BF16", 16, 1.0), ("2:4 稀疏 BF16", 9, 2.0), ("稠密 W8A16", 8, 1.0), ("稠密 W4A16", 4.125, 1.0), ("2:4 稀疏 + FP8（W8A8）", 5.0, 4.0)] for name, wbits, speed in cases: cells = [name, f"{wbits:g}"] for kind, B, n, ctx in (("decode", 1, 1, 4096), ("decode", 64, 1, 4096), ("prefill", 1, 2048, 1024.5)): st = step_time(m, chip, B, n, ctx, wbits, 16, "bf16") # 稀疏 / FP8 只让矩阵乘法变快；注意力核心（QKᵀ、PV）仍是稠密 BF16 lin = B * n * m.linear_flops_per_token() / chip.peak_flops("bf16") attn = B * n * m.attn_flops_per_token(ctx) / chip.peak_flops("bf16") t_math = lin / speed + attn t = max(t_math, st["t_mem"]) cells.append(fmt_time(t)) rows.append(cells) print_table(["方案", "比特 / 权重", "decode B=1", "decode B=64", "prefill 2048"], rows) print_observations([ "decode B = 1（读权重为主）：2:4 只把字节省到 9/16，快约 1.7 倍；同样省字节，W8A16 快 2 倍、W4A16 快 3.5 倍——" "**在 memory-bound 的 decode 里，量化比 2:4 稀疏划算**。", "prefill（compute-bound）：2:4 让矩阵乘法的算力翻倍，快约 1.9 倍（注意力核心不受益，所以不到 2 倍）；" "仅权重量化在这里一点不快（03 脚本演示 2）。", "2:4 + FP8 两种都叠上（每个权重 (2×8 + 2×2) / 4 = 5 位）：prefill 快约 3.6 倍，decode B = 1 快约 3 倍。前提是精度扛得住——演示 1 说明了稀疏的精度代价不小。" "（Blackwell Ultra 博客里 NVFP4 的稀疏算力只比稠密高 1.33 倍，第 1 讲练习 9：硬件对稀疏的支持并不总是'翻倍'。）", ]) # --------------------------------------------------------------------------- # 演示 4：激活稀疏 # --------------------------------------------------------------------------- def demo4_activation_sparsity(params, toks): print_header("演示 4：激活稀疏——batch = 1 时能跳过的，batch 一大就跳不过了") cap = [] def grab(x, W, op, layer): y = x @ W if op == "mlp_gate": cap.append(y.reshape(-1, y.shape[-1])) return y forward(params, QCFG, toks, linear_fn=grab) g = np.concatenate(cap, axis=0) # [层数 × token 数, d_ff]：W_gate 的输出（激活函数之前） relu_zero = float(np.mean(g <= 0)) silu = g / (1 + np.exp(-g)) near0 = float(np.mean(np.abs(silu) < 0.01)) print(f"MLP 门控的输入 g = x·W_gate：若用 ReLU，{relu_zero * 100:.1f}% 的神经元输出恰好为 0；" f"用 SiLU（Llama 的做法），只有 {near0 * 100:.1f}% 的输出绝对值 < 0.01，几乎没有精确的 0。") rng = np.random.default_rng(4) T = toks.size act = (g > 0).reshape(QCFG.n_layers, T, QCFG.d_ff) rows = [] for B in (1, 2, 4, 8, 16, 64): fr = [] for _ in range(50): idx = rng.choice(T, size=B, replace=False) fr.append(np.mean(np.any(act[:, idx, :], axis=1))) need = float(np.mean(fr)) hypo = 1 - 0.9 ** B rows.append([B, f"{need * 100:.1f}%", f"{(1 - need) * 100:.1f}%", f"{hypo * 100:.1f}%"]) print_table(["批大小 B", "W_down 必须读的行（本模型，ReLU）", "能跳过的", "假想：每个 token 只激活 10%"], rows) print_observations([ "ReLU 让约一半的神经元输出恰好为 0：decode 时如果事先知道哪些神经元是 0（需要一个便宜的预测器，例如 Deja Vu 的做法），" "就可以少读 W_down 的对应行、也少读 W_gate / W_up 的对应列。batch = 1 时这能省约一半的 MLP 字节。", "但 batch 里任何一个 token 用到的行都得读：B = 8 时已经要读 99.6% 的行（每个 token 独立地各激活一半，" "并集很快覆盖全部）。即使假想每个 token 只激活 10%，B = 16 时也要读 81.5%，B = 64 时 99.9%。**激活稀疏只帮得了小批量**——" "和'每个样本一个矩阵'一样，它是和批处理对着干的。", "Llama 等模型用 SiLU / SwiGLU，几乎没有精确的 0；想利用激活稀疏要先把激活函数换回 ReLU 再微调（'ReLU 化'）。", "混合专家（MoE，05 脚本）是一种**结构化、训练出来的**激活稀疏：每个 token 只用几个专家。它同样会遇到" "'批越大，用到的专家越多'的问题。", ]) if __name__ == "__main__": pa, tk, rf = setup() demo1_prune(pa, tk, rf) demo2_storage() demo3_roofline() demo4_activation_sparsity(pa, tk)