""" ============================================================================= CS349E · 前置讲 P3 · 04_softmax_and_stability.py —— 数值稳定性：softmax、log-sum-exp、在线 softmax 与方差 ============================================================================= 【这个脚本要回答的问题】 有些公式在数学上完全正确，在浮点数里却会算出 Inf、NaN 或者毫无意义的数。LLM 推理里最典型的是 softmax （注意力和最后的采样都要用）以及 LayerNorm / RMSNorm 里的方差。本脚本： 演示 1：exp 什么时候溢出？朴素 softmax 在 FP16 / FP32 里什么时候变成 NaN？减去最大值为什么能救。 演示 2：log(softmax) 与 log-sum-exp：交叉熵 / 对数概率为什么要用 log-sum-exp 来算。 演示 3：在线 softmax：一遍扫描同时求最大值和分母；分段计算再合并（FlashAttention 的数学基础）；访存次数。 演示 4：注意力分数为什么要除以 √d：不缩放时 softmax 几乎变成 one-hot，而且 FP16 下更容易溢出。 演示 5：方差的"一遍公式"E[x²] − E[x]² 的灾难性抵消；两遍法和 Welford 算法。 【怎么读输出】全部是确定性的计算，每次运行结果相同。FP16 / FP32 用的是 numpy 的原生类型（真实的 IEEE 运算）。 【运行】python3 04_softmax_and_stability.py （约 1 秒） 【依赖】numpy；从 ../lecture01 导入 console；同目录的 numfmt。 """ 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 print_header, print_observations, print_table # noqa: E402 from numfmt import (BF16, FP8_E4M3, FP16, FP32, merge_softmax_stats, online_softmax_stats, # noqa: E402 softmax_naive, softmax_safe, variance_one_pass, variance_two_pass, variance_welford) def describe(p): """把一个概率向量概括成一句话：含 NaN / 和是多少 / 最大概率。""" p = np.asarray(p, dtype=np.float64) if np.isnan(p).any(): return "NaN" return f"和 {p.sum():.4f}，最大 {p.max():.4f}" # --------------------------------------------------------------------------- # 演示 1：exp 溢出与安全 softmax # --------------------------------------------------------------------------- def demo1_exp_overflow(): """ 第一部分：每种格式里 exp(x) 不溢出的最大 x = ln(最大值)。 第二部分：把一组 logits 整体放大（乘以 scale），比较朴素 softmax 和减最大值的 softmax。 【为什么减最大值不改变结果】softmax(x)_i = e^(x_i) / Σ_j e^(x_j)。分子分母同乘 e^(−m)： = e^(x_i − m) / Σ_j e^(x_j − m)。取 m = max(x)，所有指数 ≤ 0，exp 的值在 (0, 1] 里，绝不溢出； 而且最大的那一项正好是 1，分母 ≥ 1，也不会出现 0 / 0。 【Python 小知识】np.errstate(over="ignore") 临时关掉"溢出"警告，让我们能看到 Inf / NaN 本身。 """ print_header("演示 1：exp 什么时候溢出？朴素 softmax vs 减最大值的 softmax") rows = [] for f in (FP32, BF16, FP16, FP8_E4M3): rows.append([f.name, f"{f.max_value:.4g}", f"{np.log(f.max_value):.2f}"]) print_table(["格式", "最大值", "exp(x) 不溢出的最大 x"], rows) base = np.array([3.0, 1.0, 0.2, -1.0, 2.5]) rows = [] for scale in (1, 3, 5, 30, 100): x = base * scale cells = [f"{scale} × [3, 1, 0.2, −1, 2.5]"] for dt in (np.float16, np.float32): cells.append(describe(softmax_naive(x, dt))) cells.append(describe(softmax_safe(x, dt))) rows.append(cells) print() print_table(["logits", "FP16 朴素", "FP16 减最大值", "FP32 朴素", "FP32 减最大值"], rows) print_observations([ "FP16 里只要有一个 logit 超过 11.09，exp 就溢出成 Inf；FP32 / BF16 的门槛是 88.72。FP8 E4M3 只有 6.1。", "logits 放大 5 倍（最大 15）时，FP16 朴素 softmax 已经是 NaN（Inf / Inf）；放大 30 倍（最大 90）时 FP32 也 NaN 了。" "LLM 最后一层的 logits 到几十很常见，注意力分数在不缩放时也可以很大（演示 4）。", "减去最大值以后，所有情况都正常，而且概率之和是 1——数学上一样的公式，数值上天差地别。" "所有实际的 softmax 实现（PyTorch、FlashAttention、vLLM 的采样器）都这样做。", ]) # --------------------------------------------------------------------------- # 演示 2：log-softmax 与 log-sum-exp # --------------------------------------------------------------------------- def demo2_log_softmax(): """ 对数概率 log p_i = x_i − log Σ_j e^(x_j) = x_i − LSE(x)。直接先算 softmax 再取 log： p_i 下溢成 0 时，log 0 = −Inf，交叉熵变成 Inf。用 log-sum-exp（LSE，同样先减最大值）就没有这个问题： LSE(x) = m + log Σ_j e^(x_j − m)。 【系统联系】vLLM 等推理引擎返回 logprobs、投机解码的接受率（第 6 讲）、计算困惑度，都需要对数概率； 它们都用 log_softmax（内部就是 LSE），而不是 log(softmax)。 """ print_header("演示 2：log(softmax) vs log-sum-exp——对数概率怎样算才不会变成 −Inf") rows = [] for gap in (10.0, 50.0, 120.0, 800.0): x = np.array([0.0, -gap, -gap / 2], dtype=np.float32) with np.errstate(divide="ignore", under="ignore"): naive = np.log(softmax_safe(x, np.float32)) m = x.max() lse = m + np.log(np.exp(x - m).sum(dtype=np.float32)) stable = x - lse rows.append([f"[0, −{gap:g}, −{gap / 2:g}]", f"{naive[1]:.6g}", f"{stable[1]:.6g}", f"{-gap:.6g}"]) print_table(["logits（FP32）", "log(softmax) 的第 2 项", "x − LSE(x) 的第 2 项", "精确值（≈ −差距）"], rows) print_observations([ "差距到 120 时，e^(−120) ≈ 7.7e−53 已经小于 FP32 能表示的最小正数（约 1.4e−45），softmax 给出 0，log 得到 −Inf。", "用 log-sum-exp 直接算对数概率，结果始终有限而且精确：对数空间里的数不会下溢。", "教训：能在对数空间里算的量（对数概率、对数似然、注意力的归一化常数），就不要先 exp 再 log。", ]) # --------------------------------------------------------------------------- # 演示 3：在线 softmax # --------------------------------------------------------------------------- def demo3_online_softmax(V=32_000, seed=0): """ 第一部分：长度 V 的 logits，用三种方法求 softmax 的归一化常数，比较结果： 两遍（先求最大值、再求和）、在线一遍（online_softmax_stats）、分成 8 段各自在线、再用 ⊕ 合并。 第二部分：每个元素的访存次数（Milakov & Gimelshein 2018 的数法）： 安全 softmax：读一遍求最大值、读一遍求和、读一遍算输出、写一遍输出 = 4 次； 在线 softmax：读一遍同时求最大值和和、读一遍算输出、写一遍 = 3 次。 【为什么这很重要】softmax 是 memory-bound 的（每个元素只有几次运算），访存次数直接决定时间： 论文在 V ≥ 4000 时测到最多约 1.3 倍的加速。更重要的是，在线 softmax 的 (m, d) 可以分段合并， 这让 FlashAttention 能在**不把整行分数写回 HBM** 的情况下算注意力（第 11 讲）。 """ print_header(f"演示 3：在线 softmax——一遍扫描，可以分段合并（V = {V:,} 个 logits，FP32）") rng = np.random.default_rng(seed) x = (rng.standard_normal(V) * 3).astype(np.float32) m2 = x.max() d2 = np.exp(x - m2).sum(dtype=np.float32) m1, d1 = online_softmax_stats(x) chunks = np.array_split(x, 8) stats = [online_softmax_stats(c) for c in chunks] mc, dc = stats[0] for st in stats[1:]: mc, dc = merge_softmax_stats((mc, dc), st) d_seq = np.float32(0.0) for e in np.exp(x - m2): # 两遍法，但第二遍像在线算法一样"一个一个加" d_seq = d_seq + e ref = np.exp(x.astype(np.float64) - x.max()).sum() rows = [["两遍：先 max，再 sum（numpy 的 sum）", f"{m2:.6f}", f"{d2:.6f}", f"{abs(d2 - ref) / ref:.1e}"], ["两遍：先 max，再逐个累加", f"{m2:.6f}", f"{d_seq:.6f}", f"{abs(d_seq - ref) / ref:.1e}"], ["在线，一遍（逐个累加）", f"{m1:.6f}", f"{d1:.6f}", f"{abs(d1 - ref) / ref:.1e}"], ["8 段各自在线 + ⊕ 合并", f"{mc:.6f}", f"{dc:.6f}", f"{abs(dc - ref) / ref:.1e}"]] print_table(["方法", "最大值 m", "分母 d = Σ e^(x − m)", "d 的相对误差（对 FP64）"], rows) print() rows = [["朴素（不减最大值）", "1 次（求和）+ 1 次（输出）", 1, 3, "会溢出"], ["安全（两遍）", "2 次（max、sum）+ 1 次（输出）", 1, 4, "稳定"], ["在线（Milakov & Gimelshein 2018）", "1 次（max 和 sum 一起）+ 1 次（输出）", 1, 3, "稳定"]] print_table(["softmax 算法", "读输入", "写输出", "每元素访存次数", "数值"], rows) print_observations([ f"四种方法的最大值完全相同。分母的误差差别来自**求和顺序**，而不是'在线'本身：numpy 的 sum 内部用两两求和，" f"误差 {abs(d2 - ref) / ref:.1e}；同样的两遍法改成逐个累加，误差升到 {abs(d_seq - ref) / ref:.1e}，" f"和在线一遍（也是逐个累加）的 {abs(d1 - ref) / ref:.1e} 是同一个量级（03 脚本演示 4 的结论在这里重演）。", f"分成 8 段、每段在线、再合并，误差降到 {abs(dc - ref) / ref:.1e}：一次长累加被拆成几次短累加，每段的误差都小。" "GPU 内核里每个线程只累加几个元素、再用 warp shuffle 树形合并，所以实际的在线 softmax 误差很小。", "⊕ 运算满足结合律（数学上），一行可以切给不同的线程块并行算，最后再合。" "FlashAttention（按 KV 分块）和 flash-decoding（把很长的 KV 切成几段并行）都建立在这个性质上（第 11 讲）。", "访存从每元素 4 次降到 3 次，对 memory-bound 的 softmax 就是约 1.33 倍的理论加速（论文实测最多约 1.3 倍）。" "如果再把 softmax 和前后的运算融合（例如和 top-k 融合，论文报告最多 5 倍），省得更多——这是第 11–13 讲'算子融合'的预告。", ]) # --------------------------------------------------------------------------- # 演示 4：注意力分数的缩放 # --------------------------------------------------------------------------- def demo4_attention_scaling(d=128, n_keys=4096, seed=0): """ q、k 的每个分量是标准正态，点积 q·k 的标准差是 √d。比较 softmax(q·K) 和 softmax(q·K / √d)： 最大概率、熵（越接近 log n_keys 越"均匀"）、logits 的最大值，以及 FP16 下朴素 softmax 是否溢出。 【为什么要除以 √d】（Attention Is All You Need §3.2.1）d 很大时点积的方差随 d 增长，softmax 被推到 "几乎 one-hot"的饱和区，梯度很小。除以 √d 把方差拉回 1。推理时它还有一个数值上的好处： 分数的绝对值变小，低精度下更不容易溢出。 """ print_header(f"演示 4：注意力分数为什么要除以 √d（d = {d}，{n_keys:,} 个 key，FP32 计算）") rng = np.random.default_rng(seed) q = rng.standard_normal(d) K = rng.standard_normal((n_keys, d)) rows = [] for label, s in (("不缩放：q·k", 1.0), ("除以 √d", 1 / np.sqrt(d))): z = (K @ q) * s p = softmax_safe(z, np.float64) ent = -np.sum(p * np.log(np.maximum(p, 1e-300))) naive16 = softmax_naive(z, np.float16) rows.append([label, f"{z.std():.2f}", f"{z.max():.1f}", f"{p.max():.3f}", f"{ent:.2f} / {np.log(n_keys):.2f}", "NaN" if np.isnan(naive16).any() else "正常"]) print_table(["分数", "分数的标准差", "分数的最大值", "最大注意力权重", "熵 / 均匀分布的熵", "FP16 朴素 softmax"], rows) print_observations([ f"不缩放时分数的标准差约为 √{d} ≈ {np.sqrt(d):.1f}，最大分数几十：注意力集中到极少数 key 上，熵远低于均匀分布；" "而且在 FP16 里朴素 softmax 直接溢出。", "除以 √d 后标准差约为 1，注意力分布正常。FlashAttention 等内核通常把 1/√d 和 log2(e) 合并成一个常数，" "乘到分数上再用 exp2 计算（exp2 在硬件上比 exp 便宜）。", "第 1 讲 §5.5 讲过 softmax 的 exp 会成为 Blackwell 上的瓶颈；这里看到的是它的另一面：exp 的数值范围也要小心。", ]) # --------------------------------------------------------------------------- # 演示 5：方差的灾难性抵消 # --------------------------------------------------------------------------- def demo5_variance(n=4096, seed=0): """ 均值很大、方差很小的数据（例如 mean = 1e4、std = 1），在 FP32 里用三种方法算方差。 【灾难性抵消 catastrophic cancellation】E[x²] ≈ 1e8，E[x]² ≈ 1e8，两者相减得到 1。但 1e8 在 FP32 里的 ulp 是 8： 两个大数各自带着几个单位的舍入误差，相减后"有效数字"全被抵消掉了，剩下的主要是误差。 【两遍法】先求均值，再求 Σ(x − mean)²：x − mean 都是 O(1) 的小数，没有抵消。代价是读两遍数据。 【Welford】一遍扫描、增量更新均值和离差平方和，数值上和两遍法一样稳定。 """ print_header(f"演示 5：方差——一遍公式 E[x²] − E[x]² 的灾难性抵消（n = {n:,}，FP32）") rng = np.random.default_rng(seed) rows = [] for mean in (0.0, 1e2, 1e3, 1e4): x = (mean + rng.standard_normal(n)).astype(np.float32) ref = np.var(x.astype(np.float64)) v1 = variance_one_pass(x) v2 = variance_two_pass(x) vw = variance_welford(x) rows.append([f"{mean:g}", f"{ref:.6f}", f"{v1:.6f}", f"{v2:.6f}", f"{vw:.6f}"]) print_table(["均值（标准差 = 1）", "精确方差（FP64）", "一遍公式（FP32）", "两遍法（FP32）", "Welford（FP32）"], rows) print_observations([ "均值为 0 时三种方法都对；均值到 1e3 时一遍公式已经明显偏离，到 1e4 时结果完全错误——方差竟然是负数。", "两遍法和 FP64 几乎一模一样；Welford 的误差略大（均值 1e4 时从小数点后第 4 位开始不同），但都远好于一遍公式。", "LayerNorm 要先算均值和方差，激活里如果有很大的常数分量（离群通道，第 3 讲的 LLM.int8 会讲到），" "一遍公式就会出问题。融合的 LayerNorm 内核想只读一遍数据，就用 Welford（以及它的分块合并版本，" "和演示 3 的 ⊕ 是同一个思路）。RMSNorm 只需要 E[x²]，没有这个抵消问题——这也是它更受欢迎的原因之一。", ]) if __name__ == "__main__": demo1_exp_overflow() demo2_log_softmax() demo3_online_softmax() demo4_attention_scaling() demo5_variance()