""" ============================================================================= 第 5 讲 · 代码 3：让内存受限的运算变快 —— 低精度、算子融合、重计算 ============================================================================= 对应讲义：lectures/lecture05_gpus_tpus.md 第 4 节（技巧 2、3、4） ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 01 号文件告诉我们：逐元素运算、归一化这类"低算术强度"的运算完全被内存带宽卡住。 既然瓶颈是"搬数据"，那就想办法**少搬**： 技巧 1：低精度 —— 每个数用更少的字节，搬运量直接减半 技巧 2：算子融合 —— 把多个运算合成一个核函数，中间结果留在芯片上，不写回显存 技巧 3：重计算 —— 反向传播需要的中间结果不存了，需要时重新算（用计算换搬运） 每个技巧都先用"数一数读写次数"的模型讲清楚，再在你的电脑上实测或验证。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 03_fusion_precision_recompute.py （首次运行 torch.compile 需要编译，约 10~30 秒） """ import time import torch def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def get_device(): """优先使用 GPU（CUDA / Apple MPS），否则用 CPU。""" if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def sync(device): """等待设备上的运算完成。""" if device.type == "cuda": torch.cuda.synchronize() elif device.type == "mps": torch.mps.synchronize() def bench(fn, device, repeats=20): """热身后计时，返回每次平均秒数。""" for _ in range(3): fn() sync(device) t0 = time.perf_counter() for _ in range(repeats): fn() sync(device) return (time.perf_counter() - t0) / repeats # ============================================================================= # 第 1 部分：低精度 # ============================================================================= def demo_low_precision(device): """演示 1：同样多的元素，BF16 的搬运时间是 FP32 的一半。""" print_header("演示 1：低精度 —— 字节少了，搬得就快") n = 1 << 25 for dt in [torch.float32, torch.bfloat16]: x = torch.randn(n, device=device).to(dt) t = bench(lambda: x.clone(), device) print(f" 复制 {n:,} 个 {str(dt):<15}：{t * 1e3:.2f} ms（{2 * n * x.element_size() / t / 1e9:.0f} GB/s）") print(""" 带宽（GB/s）几乎一样，但 BF16 的字节数只有一半，所以时间也只要一半。 对内存受限的运算，精度减半 ≈ 速度翻倍；对矩阵乘法，GPU 的 Tensor Core 在低精度下的算力也高得多 （H100：BF16 约 990 TF/s，FP32 约 67 TF/s）。这是第 2 讲混合精度训练的另一个重要动机。 以 ReLU 为例（CS336 课件）：FP32 时每次运算要搬 8 字节，FP16/BF16 只要 4 字节 —— 算术强度翻倍。""") # ============================================================================= # 第 2 部分：算子融合 # ============================================================================= def sin2_plus_cos2(x): """sin²(x) + cos²(x)（结果恒为 1，CS336 课件的例子）。朴素执行时会启动 5 个核函数。""" return torch.sin(x) ** 2 + torch.cos(x) ** 2 def gelu_tanh(x): """GELU 的 tanh 近似，用基本运算手写（每个运算都是一个单独的核函数）。""" return 0.5 * x * (1.0 + torch.tanh(0.7978845608 * (x + 0.044715 * x ** 3))) def memory_traffic_sin_cos(n): """ 数一数 sin²(x) + cos²(x) 的显存读写（以"元素"为单位）。 【未融合】5 个核函数，每个都要从显存读输入、把结果写回显存： a = sin(x) 读 x 写 a b = a² 读 a 写 b c = cos(x) 读 x 写 c d = c² 读 c 写 d e = b + d 读 b、d 写 e 共读 6n、写 5n 【融合】一个核函数：读 x，在寄存器里算完全部 5 步，写 e。共读 n、写 n。 """ return {"未融合": (6 * n, 5 * n), "融合": (n, n)} def demo_fusion(device): """演示 2：算子融合 —— 模型与实测。""" print_header("演示 2：算子融合") n = 1 << 24 for name, (r, w) in memory_traffic_sin_cos(n).items(): print(f" sin²+cos²，{name}：读 {r // n}n、写 {w // n}n 个元素，共 {(r + w) // n}n") print(" → 融合后搬运量减少到 2/11，理论上可以快约 5.5 倍（在内存受限的前提下）。") x = torch.randn(n, device=device) print(f"\n 在你的电脑上实测（设备：{device}，{n:,} 个元素）：") try: # 【PyTorch 小知识】torch.compile(函数) 会追踪函数里的运算，把能合并的逐元素运算融合成一个核函数 # （在 NVIDIA GPU 上生成 Triton 核函数，在 CPU 上生成 C++ 代码）。第一次调用时编译，之后直接用。 compiled = {f.__name__: torch.compile(f) for f in [sin2_plus_cos2, gelu_tanh]} for f in [sin2_plus_cos2, gelu_tanh]: cf = compiled[f.__name__] ref = f(x) out = cf(x) ok = torch.allclose(ref, out, atol=1e-5) t_eager = bench(lambda: f(x), device) t_comp = bench(lambda: cf(x), device) print(f" {f.__name__:<16} 逐个运算 {t_eager * 1e3:6.2f} ms，torch.compile 融合后 {t_comp * 1e3:6.2f} ms，" f"加速 {t_eager / t_comp:.1f} 倍（结果一致：{ok}）") except Exception as e: # 某些环境下 torch.compile 不可用（比如缺少编译器） print(f" torch.compile 在这台电脑上不可用（{type(e).__name__}），跳过实测。") t_builtin = bench(lambda: torch.nn.functional.gelu(x, approximate="tanh"), device) print(f" 对比：PyTorch 内置的 F.gelu（本身就是一个融合好的核函数）{t_builtin * 1e3:.2f} ms") print(""" 在 GPU 上，融合通常能带来数倍的加速：几乎所有时间都花在读写显存上，而融合把 11 次读写减少到 2 次。 在 CPU 上加速往往没那么明显：sin、cos、tanh 这类运算在 CPU 上计算本身就比较慢，不完全是内存受限。 简单的逐元素融合 torch.compile 可以自动完成；更复杂的融合（如 FlashAttention）需要手写核函数（第 6 讲）。""") # ============================================================================= # 第 3 部分：重计算 # ============================================================================= class ThreeSigmoidsStore(torch.autograd.Function): """ y = σ(σ(σ(x)))，**保存**所有中间结果供反向使用（普通做法）。 【PyTorch 小知识】torch.autograd.Function 让你自己定义前向和反向： forward(ctx, ...) ：计算输出，用 ctx.save_for_backward(...) 保存反向需要的张量 backward(ctx, grad)：拿到"损失对输出的梯度"，返回"损失对输入的梯度" σ 的导数：σ'(u) = σ(u)(1 - σ(u))，只需要 σ(u) 本身就能算出来。 """ @staticmethod def forward(ctx, x): s1 = torch.sigmoid(x) s2 = torch.sigmoid(s1) s3 = torch.sigmoid(s2) ctx.save_for_backward(s1, s2, s3) # 存 3 个中间结果 return s3 @staticmethod def backward(ctx, grad_out): s1, s2, s3 = ctx.saved_tensors g = grad_out * s3 * (1 - s3) g = g * s2 * (1 - s2) return g * s1 * (1 - s1) class ThreeSigmoidsRecompute(torch.autograd.Function): """ 同样的函数，但只保存输入 x，反向时**重新计算** s1、s2、s3。 【为什么可能更快】在一个融合的核函数里，重算 3 个 sigmoid 只是寄存器里的几次运算， 而"把 3 个中间结果写回显存、反向时再读出来"要搬运 6n 个元素。对内存受限的运算， 多算几次比多搬几次便宜。这就是 FlashAttention 反向传播的思路（重算注意力矩阵而不是存它）。 """ @staticmethod def forward(ctx, x): ctx.save_for_backward(x) # 只存输入 return torch.sigmoid(torch.sigmoid(torch.sigmoid(x))) @staticmethod def backward(ctx, grad_out): (x,) = ctx.saved_tensors s1 = torch.sigmoid(x) # 重新算一遍前向 s2 = torch.sigmoid(s1) s3 = torch.sigmoid(s2) g = grad_out * s3 * (1 - s3) g = g * s2 * (1 - s2) return g * s1 * (1 - s1) def saved_bytes(fn): """测量执行 fn() 时为反向传播保存的字节数（方法见第 2 讲 05 号文件）。""" seen = {} def pack(t): seen[t.data_ptr()] = t.numel() * t.element_size() return t with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t): out = fn() return sum(seen.values()), out def demo_recompute(): """演示 3：重计算 —— 少存、少搬，梯度完全相同。""" print_header("演示 3：重计算（以三个 sigmoid 叠加为例）") n = 1 << 20 x1 = torch.randn(n, requires_grad=True) x2 = x1.detach().clone().requires_grad_(True) b1, y1 = saved_bytes(lambda: ThreeSigmoidsStore.apply(x1)) b2, y2 = saved_bytes(lambda: ThreeSigmoidsRecompute.apply(x2)) y1.sum().backward() y2.sum().backward() print(f" 保存中间结果：为反向保存了 {b1 / 1e6:.1f} MB；重计算：只保存了 {b2 / 1e6:.1f} MB") print(f" 两种方式的梯度最大差异：{(x1.grad - x2.grad).abs().max().item():.1e}") ok = torch.autograd.gradcheck(ThreeSigmoidsRecompute.apply, (torch.randn(8, dtype=torch.float64, requires_grad=True),)) print(f" torch.autograd.gradcheck（用有限差分检查反向是否写对）：{ok}") print(""" 显存读写（以 n 个元素为单位；假设每个 Function 内部是融合的核函数）： 前向 反向 合计 保存中间结果：读 x，写 s1、s2、s3（1+3） 读 上游梯度、s1、s2、s3，写 输入梯度（4+1） 9 重计算 ：读 x，写 s3（1+1） 读 上游梯度、x，写 输入梯度（2+1） 5 重计算多做了 3 次 sigmoid，却少搬了近一半的数据。对内存受限的运算，这笔交易很划算。 （第 2 讲的激活检查点是同样的思想，只是粒度是"整层"而不是"一个融合核函数内部"。）""") def main(): torch.manual_seed(0) device = get_device() demo_low_precision(device) demo_fusion(device) demo_recompute() print("\n全部演示运行完毕。") if __name__ == "__main__": main()