18c · 剪枝与稀疏化

import time, torch

torch.manual_seed(0)
torch.set_num_threads(4)
W = torch.randn(4096, 4096)

k = int(W.numel() * 0.90)
thr = W.abs().flatten().kthvalue(k).values
W_pruned = W * (W.abs() > thr)          # ⭐ 幅度剪枝：|w| 最小的 90% 置零

print("稀疏度        : %.1f%%" % ((W_pruned == 0).float().mean() * 100))
print("原张量占用    : %.1f MiB" % (W.element_size() * W.nelement() / 2**20))
print("剪枝后占用    : %.1f MiB" % (W_pruned.element_size() * W_pruned.nelement() / 2**20))
# ⭐ 上面两行一定相等：置零不改变形状，dtype 也没变，显存一个字节都没省

W_csr = W_pruned.to_sparse_csr()        # 换成稀疏格式才可能省，但索引本身也要钱
csr_bytes = (W_csr.values().numel() * 4
             + W_csr.col_indices().numel() * 8
             + W_csr.crow_indices().numel() * 8)
print("CSR 占用      : %.1f MiB" % (csr_bytes / 2**20))
# 该看到的数字（4096×4096 FP32，90% 稀疏度，已实跑）：
#   原张量 64.0 MiB → 剪枝后 64.0 MiB（一个字节没省）
#   → CSR 19.2 MiB  ⭐ 而"只存非零值"的理想值是 6.4 MiB
#   多出来的 12.8 MiB 全是行/列索引 —— 稀疏格式的账要连索引一起算

def bench(fn, warmup=3, runs=10):
    for _ in range(warmup):
        fn()
    t0 = time.perf_counter()
    for _ in range(runs):
        fn()
    return (time.perf_counter() - t0) / runs * 1000

for B in [1, 64, 1024]:                 # B=1 就是 LLM decode 的形状
    x = torch.randn(B, 4096)
    td = bench(lambda: x @ W_pruned.T)
    ts = bench(lambda: torch.sparse.mm(W_csr, x.T))
    print(f"B={B:<5} 稠密 {td:8.2f} ms   稀疏 {ts:8.2f} ms   ⭐ 比值 {ts/td:.2f}x")

   稀疏度        : 90.0%
   原张量占用    : 64.0 MiB
   剪枝后占用    : 64.0 MiB      ⭐ 一个字节没省
   CSR 占用      : 19.2 MiB      （值 25.6MB 的 1/4，但索引又吃掉一大块）

   B=1     稠密   1.61 ms   稀疏   1.76 ms   ⭐ 比值 1.09x  ← 反而更慢！
   B=64    稠密  14.56 ms   稀疏   5.26 ms      比值 0.36x
   B=1024  稠密 366.22 ms   稀疏 215.33 ms      比值 0.59x

   ⭐ 90% 的乘法被省掉了，但最好的情况也只快了不到 3 倍，
     而在 B=1（也就是【自回归解码的形状】）上【反而更慢】💀

   稠密 matmul 为什么快：
   · 权重在显存里【连续】→ 访存可以合并（coalesced），带宽跑满
   · 形状规整 → 直接进 Tensor Core

   稀疏 matmul 为什么不快：
   · 每个非零值旁边还得读一个【索引】→ 多搬了 50~100% 的字节 ⚠️
   · 索引指向的位置是【随机的】→ 访存合并失效，实际带宽掉一大截
   · 形状不规整 → Tensor Core 用不上，退回普通 CUDA core 💀

   ⭐ 最关键的一条：
     剪枝省掉的是【乘法次数】= 算力，
     而解码本来就是【带宽瓶颈】，算力有 95% 在闲置（第 15 章）——
     你省掉了一个本来就不是瓶颈的东西，等于没省。

   规则：每【连续 4 个权重】里，恰好有 2 个是 0

   [ w0  w1  w2  w3 ] → [ 0  w1  0  w3 ]   ✅ 合法
                      → [ 0   0  0  w3 ]   ❌ 不合法（3 个零）

   ⭐ Ampere（A100）起，稀疏 Tensor Core 原生支持这个模式
      → 理论 2 倍算力，实测 1.3~1.7 倍
   ⚠️ 但稀疏度被【锁死在 50%】，压缩比远不如非结构化

   幅度剪枝（magnitude pruning）：|w| 越小越先剪
   → 极简单、极强的 baseline，至今没被大幅超越

   ⚠️ 全局阈值 vs 逐层阈值：
     全局排序容易把某一层【整层剪光】（各层权重尺度天然不同）
     → 实践上按层设配额更稳

   ⭐ 一次剪到 90% ≠ 分 5 次剪到 90%

   一次到位：  90% ────────────▶ 精度崩
   迭代：      50% → 微调 → 70% → 微调 → 85% → 微调 → 90%  ✅
                                        每次让网络有机会重新分配容量

   预期                    实测
   显存 26GB → 2.6GB      💀 26.1GB（不降反升，mask 本身还占了 0.1GB）
   P50 decode 41ms        💀 43ms（略微变慢）

   ⭐ ① 压缩项的验收指标只有两个：【显存快照】和【P50/P99 延迟】
        「稀疏度」「参数量」「压缩比」全是过程指标，不能当验收
   ⭐ ② 立项当天先做 30 分钟的「形状体检」（本节那段代码）：
        在【目标硬件】上确认这种稀疏到底能不能被吃掉
      ③ 要真加速就换路线：结构化剪枝，或者 2:4 半结构化