""" ============================================================================= CS349E · 第 11 讲 · 03_flash_backward.py —— 反向传播：重算比存储便宜；causal 跳过一半的块 ============================================================================= 【这个脚本要回答的问题】 标准注意力的反向传播要用前向的 P（N×N）。FlashAttention 不存 P，只存每行一个 logsumexp L， 反向时在片上重算 S、P——多做了运算，却更快、更省显存。 演示 1：梯度检查：FlashAttention 的反向 = 标准反向 = 有限差分。 演示 2：显存：标准写法为反向存下的 P vs FlashAttention 存的 L。 演示 3：前向 + 反向的运算量和读写随 N 的变化，H100 上的 roofline 时间。 演示 4：因果掩码：整块被掩码的 (i, j) 直接跳过，运算和读写各省多少。 【运行】python3 03_flash_backward.py （约 5 秒） 【依赖】numpy；本目录的 flash_lab；../lecture01 的 console。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from flash_lab import flash_bwd, flash_fwd, standard_attention, standard_attention_bwd # noqa: E402 from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 def std_fwd_bwd_io(N, d): """标准注意力前向 + 反向的读写元素数和矩阵乘法 FLOPs（和 flash_lab 里两个函数的计数公式相同，不实际计算 N×N 矩阵）。""" fwd = (2 * N * d + 2 * N * N + N * d) + (2 * N * N + N * d) bwd = (5 * N * N + 2 * N * d + N * d + N * d + N * d + N * d + N) + (2 * N * N + N * d + N * d + N * d + N) flops = 4 * N * N * d + 8 * N * N * d return fwd + bwd, flops def demo1_gradcheck(): print_header("演示 1：梯度检查（N = 48，d = 16，块 16×12，FP64）") rng = np.random.default_rng(0) N, d = 48, 16 Q, K, V = (rng.standard_normal((N, d)) for _ in range(3)) G = rng.standard_normal((N, d)) # 损失 = Σ O ∘ G，所以 dO = G rows = [] for causal in (False, True): O, P, _ = standard_attention(Q, K, V, causal) dQs, dKs, dVs, _ = standard_attention_bwd(Q, K, V, P, G) O2, L, _ = flash_fwd(Q, K, V, 16, 12, causal) dQ, dK, dV, _ = flash_bwd(Q, K, V, O2, G, L, 16, 12, causal) eps = 1e-6 num = {} for name, X in (("Q", Q), ("K", K), ("V", V)): idx = [(int(rng.integers(N)), int(rng.integers(d))) for _ in range(6)] vals = [] for (r, c) in idx: Xp, Xm = X.copy(), X.copy() Xp[r, c] += eps Xm[r, c] -= eps args_p = {"Q": Q, "K": K, "V": V} args_m = dict(args_p) args_p[name], args_m[name] = Xp, Xm fp = np.sum(standard_attention(args_p["Q"], args_p["K"], args_p["V"], causal)[0] * G) fm = np.sum(standard_attention(args_m["Q"], args_m["K"], args_m["V"], causal)[0] * G) vals.append(((fp - fm) / (2 * eps), r, c)) num[name] = vals grads = {"Q": dQ, "K": dK, "V": dV} fd_err = max(abs(v - grads[n][r, c]) for n in num for v, r, c in num[n]) rows.append(["是" if causal else "否", f"{np.abs(O2 - O).max():.1e}", f"{max(np.abs(dQ - dQs).max(), np.abs(dK - dKs).max(), np.abs(dV - dVs).max()):.1e}", f"{fd_err:.1e}"]) print_table(["causal", "前向 O 的差", "FlashAttention 反向 vs 标准反向（dQ、dK、dV 最大差）", "vs 有限差分（18 个随机元素）"], rows) print_observations([ "FlashAttention 的反向只用了 Q、K、V、O、dO 和每行一个 L = m + log l：P_ij = exp(S_ij − L_i) 在片上重算出来，" "和标准写法存下的 P 完全一样，所以梯度也一样（1e-15 量级）。", "有限差分的误差约 1e-9（中心差分的截断误差 + 舍入），说明两种反向的公式本身是对的。", "关键的一步是 dS = P ∘ (dP − D)，D_i = Σ_j dP_ij P_ij。直接算 D 要用整行 P；FlashAttention 用恒等式 D_i = Σ dO_i ∘ O_i，" "只需要 O 和 dO 的一行——这样 D 可以在反向开始前一次算好（讲义 §5.1）。", ]) def demo2_memory(): print_header("演示 2：为了反向传播要存什么（每层；批 × 头 = 32 个注意力矩阵，FP16）") rows = [] for N in (2048, 8192, 32768, 131072): p_bytes = 32 * N * N * 2 l_bytes = 32 * N * 4 rows.append([N, fmt_bytes(p_bytes), fmt_bytes(l_bytes), f"{p_bytes / l_bytes:,.0f}"]) print_table(["序列长度 N", "标准：存 P（N×N）", "FlashAttention：存 L（N 个 FP32）", "倍数"], rows) print_observations([ "标准写法为反向存下的 P 随 N² 增长：N = 32K 时每层 64 GiB——一块 80 GB 的 GPU 连一层都存不下。" "FlashAttention 只存 L，随 N 线性增长（论文定理 1：除输入输出外只需 O(N) 的额外显存）。", "这是长上下文训练能够进行的前提；推理时没有反向传播，但 prefill 的显存同样受益：不需要 N×N 的临时矩阵。", ]) def demo3_fwd_bwd_cost(): print_header("演示 3：前向 + 反向的运算量与读写（一个头，d = 128，块 128×128，所有读写都算成显存），H100：989 TFLOPS、3.35 TB/s") rng = np.random.default_rng(1) d = 128 rows = [] for N in (512, 1024, 2048, 4096, 8192): Q, K, V, dO = (rng.standard_normal((N, d)) for _ in range(4)) O, L, f = flash_fwd(Q, K, V, 128, 128) _, _, _, b = flash_bwd(Q, K, V, O, dO, L, 128, 128) s_io, s_fl = std_fwd_bwd_io(N, d) fa_io, fa_fl = f.total + b.total, f.flops + b.flops t_s = max(s_fl / 989e12, s_io * 2 / 3.35e12) t_f = max(fa_fl / 989e12, fa_io * 2 / 3.35e12) rows.append([N, f"{fa_fl / s_fl:.3f}", fmt_bytes(s_io * 2), fmt_bytes(fa_io * 2), f"{s_io / fa_io:.1f}", fmt_time(t_s), fmt_time(t_f), f"{t_s / t_f:.1f}"]) print_table(["N", "FLOPs：FA / 标准", "读写：标准", "读写：FA", "标准 / FA", "标准（roofline）", "FA（roofline）", "加速"], rows) print_observations([ "FlashAttention 前向 + 反向的矩阵乘法运算是标准写法的 7/6 ≈ 1.167 倍：前向 2 次矩阵乘法、反向标准写法 4 次、FlashAttention 5 次（多一次重算 S）。", "读写只少了约 1.8 倍，而且不随 N 变：FlashAttention 的反向对每个 (i, j) 块都要读 Q_i、dO_i 并读写 dQ_i，" "d = 128、块 128×128 时这部分约 4N² 个元素，和标准写法反复读写的 N×N 矩阵（约 11N²）同一个量级。" "按 roofline 两者都受限于带宽，FlashAttention 快约 1.8 倍。", "论文和实际测到的加速更大，我们的计数漏掉了两件事：(1) 标准实现还有单独的掩码、softmax、dropout 内核（02 脚本演示 4）；" "(2) 被反复读的 Q_i、dO_i、dQ_i 块在很多 CTA 之间共享，大部分在 L2 里命中、不到显存——我们把所有读写都算成了显存。" "这一项要用真机的 profiler 才能确定。", "\"重算比存储便宜\"是一个反复出现的系统原则：激活值重算（gradient checkpointing）、KV cache 的重算 vs 换出（第 7 讲）、" "megakernel 里 RMSNorm 的重复计算（05 脚本）都是同一个权衡——算力便宜、带宽和容量贵。", ]) def demo4_causal(): print_header("演示 4：因果掩码——整块被掩码的 (i, j) 跳过（N = 4096，d = 64，块 Br × Bc）") rng = np.random.default_rng(2) N, d = 4096, 64 Q, K, V = (rng.standard_normal((N, d)) for _ in range(3)) rows = [] for Br, Bc in ((64, 64), (128, 128), (256, 256), (128, 64)): _, _, full = flash_fwd(Q, K, V, Br, Bc, causal=False) _, _, cz = flash_fwd(Q, K, V, Br, Bc, causal=True) T = N // Bc rows.append([f"{Br}×{Bc}", f"{cz.flops / full.flops:.3f}", f"{0.5 + 0.5 / T:.3f}" if Br == Bc else "", f"{cz.total / full.total:.3f}"]) print_table(["块", "FLOPs：causal / 不带掩码", "公式 1/2 + 1/(2T)（方块时）", "读写：causal / 不带掩码"], rows) print_observations([ "因果掩码让一半的分数矩阵是 −∞。整块都被掩码的 (i, j) 不读、不算；只有对角线上的块要在块内加掩码。" "T = N/B 个块行时，算的块数是 T(T+1)/2，占比 1/2 + 1/(2T)，块越小越接近一半。", "对角线上的块里仍有约一半是被掩码掉的（算了但不用），所以块越大浪费越多——这是选块大小的又一个因素。", "读写也几乎减半：被跳过的块的 K、V 不用读。对 decode（1 个新 token 看所有缓存）没有影响：所有 key 都可见。", ]) if __name__ == "__main__": demo1_gradcheck() demo2_memory() demo3_fwd_bwd_cost() demo4_causal()