""" ============================================================================= 第 6 讲 · 代码 1：基准测试与性能分析 —— 先测量，再优化 ============================================================================= 对应讲义：lectures/lecture06_kernels_triton.md 第 2、3 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 优化代码的第一步不是写核函数，而是**测量**：到底慢在哪里？优化之后到底快了多少？ 两种工具： - 基准测试（benchmarking）：测一段代码从头到尾花了多少时间（端到端）。 它告诉你"多快"，但不告诉你"时间花在哪"。 - 性能分析（profiling）：列出执行过程中每个运算 / 核函数各花了多少时间、调用了几次。 它告诉你"时间花在哪"。 GPU 上的测量有两个大坑，这个文件会亲手演示： 1. 不热身：第一次运行包含编译、内存分配等一次性开销，会严重高估耗时； 2. 不同步：GPU 是异步执行的，Python 发出命令后立刻返回。不等 GPU 算完就停表， 测到的只是"发命令"的时间，可能比真实耗时小几百倍。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 01_benchmarking_profiling.py （有 NVIDIA GPU 或 Apple 芯片 GPU 时会演示异步的坑；纯 CPU 电脑上会跳过那一部分。） """ import statistics import time import torch from torch.profiler import ProfilerActivity, profile 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 synchronize(device): """等待设备上已经发出的所有运算完成。CPU 是同步执行的，不需要等。""" if device.type == "cuda": torch.cuda.synchronize() elif device.type == "mps": torch.mps.synchronize() # ============================================================================= # 第 1 部分：正确的基准测试函数 # ============================================================================= def benchmark(fn, device, num_warmups=3, num_trials=10): """ 正确地测量 fn() 的耗时（秒），返回每次试验的耗时列表。 【步骤】（与 CS336 课上的 benchmark 函数相同） 1. 热身 num_warmups 次：让一次性的开销（编译核函数、分配显存、选择算法）先发生； 2. 同步：确保热身的运算都已完成，不会算进后面的计时； 3. 每次试验：开始计时 → 运行 → **同步** → 停止计时； 4. 返回所有试验的耗时（可以看平均值、中位数、波动）。 【CUDA 上的更精确做法】 用 torch.cuda.Event(enable_timing=True) 在 GPU 的时间线上打点： start.record(); fn(); end.record(); torch.cuda.synchronize(); start.elapsed_time(end) 这样测到的是 GPU 上的执行时间，不受 CPU 端 Python 开销的影响。这里为了通用，用 CPU 计时 + 同步。 【参数】 fn : 无参数的函数 device : torch.device num_warmups : 热身次数（>= 0） num_trials : 试验次数（>= 1） """ for _ in range(num_warmups): fn() synchronize(device) times = [] for _ in range(num_trials): t0 = time.perf_counter() fn() synchronize(device) times.append(time.perf_counter() - t0) return times def demo_pitfalls(): """演示 1：两个大坑 —— 不热身、不同步。""" device = get_device() print_header(f"演示 1：基准测试的两个大坑（设备：{device}）") a = torch.randn(4096, 4096, device=device) b = torch.randn(4096, 4096, device=device) # 坑 1：第一次运行 t0 = time.perf_counter() a @ b synchronize(device) first = time.perf_counter() - t0 steady = statistics.mean(benchmark(lambda: a @ b, device)) print(f" 第一次运行：{first * 1e3:7.1f} ms；热身之后的平均：{steady * 1e3:7.1f} ms（第一次慢了 {first / steady:.1f} 倍）") # 坑 2：不同步 if device.type == "cpu": print(" （CPU 是同步执行的，不存在'不同步'的坑；在 GPU 上运行本文件可以看到这个现象。）") else: synchronize(device) t0 = time.perf_counter() for _ in range(10): a @ b no_sync = (time.perf_counter() - t0) / 10 # 没有等 GPU 算完就停表 synchronize(device) with_sync = (time.perf_counter() - t0) / 10 # 等 GPU 真正算完 print(f" 不同步就停表：每次 {no_sync * 1e3:7.3f} ms ← 错误！只是'发命令'的时间") print(f" 同步之后停表：每次 {with_sync * 1e3:7.3f} ms （相差约 {with_sync / no_sync:.0f} 倍）") print(""" 结论：GPU 计时必须（1）先热身，（2）在停表前同步。忘了任何一条，结论都可能完全错误。 同理，PyTorch 代码里偶尔出现的 .item()、print(张量)、.cpu() 都会隐式地同步 —— 训练循环中每一步都打印损失，会让 CPU 每一步都等 GPU，失去异步执行带来的重叠。""") def demo_scaling(): """演示 2：矩阵乘法的耗时随尺寸的变化。""" device = get_device() print_header("演示 2：耗时随尺寸的变化 —— 小时看开销，大时看 n³") print(f" {'n':>6}{'耗时':>12}{'相对 n=256':>12}{'n³ 之比':>10}") base = None for n in [64, 128, 256, 512, 1024, 2048, 4096]: a = torch.randn(n, n, device=device) b = torch.randn(n, n, device=device) t = statistics.median(benchmark(lambda: a @ b, device, num_trials=5)) if n == 256: base = t print(f" {n:>6}{t * 1e3:>10.3f} ms" + (f"{t / base:>12.1f}{(n / 256) ** 3:>10.1f}" if base else "")) print(""" 小尺寸时耗时几乎不变：固定开销（启动核函数、Python 调用）占主导，硬件远没有被喂饱； 尺寸足够大后，耗时大约按 n³ 增长（与 FLOPs 成正比）。这和 CS336 课上在 B200 上测到的规律一致。""") # ============================================================================= # 第 2 部分：性能分析 # ============================================================================= def manual_gelu(x): """用基本运算手写的 GeLU（tanh 近似）。""" return 0.5 * x * (1 + torch.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x))) def profile_ops(fn, top=8): """ 用 PyTorch 性能分析器运行 fn()，打印调用次数最多、耗时最多的运算。 【PyTorch 小知识】 with profile(activities=[ProfilerActivity.CPU]) as prof: 运行代码 prof.key_averages() 按运算名字汇总：每个运算调用了几次、总共花了多少时间。 在 NVIDIA GPU 上加入 ProfilerActivity.CUDA，就能看到每个 GPU 核函数的名字和耗时。 """ fn() # 热身 with profile(activities=[ProfilerActivity.CPU]) as prof: fn() events = sorted(prof.key_averages(), key=lambda e: -e.self_cpu_time_total) for e in events[:top]: print(f" {e.key:<28} 调用 {e.count:>3} 次，自身耗时 {e.self_cpu_time_total / 1e3:7.3f} ms") def demo_profiler(): """演示 3：性能分析器 —— 手写 GeLU vs 内置 GeLU。""" print_header("演示 3：性能分析器 —— 时间花在哪里？") x = torch.randn(1 << 22) print(" 手写的 GeLU（每个运算都是一个核函数）：") profile_ops(lambda: manual_gelu(x)) print(" PyTorch 内置的 F.gelu（一个融合好的核函数）：") profile_ops(lambda: torch.nn.functional.gelu(x, approximate="tanh")) print(""" 手写版本调用了 6 次 aten::mul、2 次 aten::add、1 次 aten::tanh：9 个核函数，每个都要读写整个张量； 内置版本只有 1 个 aten::gelu。（表里的 aten::to / aten::_to_copy 等大概率是把 Python 数字常量转换成张量的开销，耗时很小。） 在 NVIDIA GPU 上，性能分析器还会显示真正执行的 GPU 核函数名字，例如 "cutlass3x_sm100_simt_sgemm_f32_..._64x64x16" 这样的名字告诉你：用的是 CUTLASS 库、B200（sm100）架构、 FP32、64×64×16 的分块（CS336 课上的例子）。 更深入的分析可以用 NVIDIA Nsight Systems（作业 2 会用到），它能显示 CPU 和 GPU 的完整时间线。""") def main(): torch.manual_seed(0) demo_pitfalls() demo_scaling() demo_profiler() print("\n全部演示运行完毕。") if __name__ == "__main__": main()