""" ============================================================================= 第 6 讲 · 代码 6：torch.compile —— 编译器能自动做什么，不能做什么 ============================================================================= 对应讲义：lectures/lecture06_kernels_triton.md 第 10 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 手写核函数很强大，但也很费事。PyTorch 2 提供了 torch.compile：一行代码， 让编译器自动把你的 PyTorch 代码变成融合的核函数（在 NVIDIA GPU 上生成的正是 Triton 代码）。 它由两部分组成： - TorchDynamo：在 Python 执行时"截获"PyTorch 运算，记录成一张计算图； - TorchInductor：对计算图做优化（主要是融合），生成 Triton（GPU）或 C++（CPU）代码。 这个文件演示： 1. 编译一个"RMSNorm + 前馈层 + 残差"的小模块：结果一致、速度如何； 2. "图中断"（graph break）：Dynamo 遇到它无法追踪的 Python 代码时会把图切断，融合机会随之丢失； 3. 重新编译：输入形状变化时，编译器可能要重新编译； 4. 它做不到的事：需要改变算法的优化（比如 FlashAttention 的在线 softmax），还得靠人写核函数。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 06_torch_compile.py （首次编译需要十几秒到几十秒） """ import logging import statistics import time import torch import torch._dynamo as dynamo from torch._dynamo.utils import counters # 编译器会打印很多警告日志（例如"这里发生了图中断"）。本文件会用自己的话解释这些现象， # 所以把它们的日志级别调高，只显示真正的错误。想看原始日志时，把下面两行注释掉即可。 logging.getLogger("torch._dynamo").setLevel(logging.ERROR) logging.getLogger("torch._inductor").setLevel(logging.ERROR) 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): """等待设备上的运算完成。""" if device.type == "cuda": torch.cuda.synchronize() elif device.type == "mps": torch.mps.synchronize() def bench(fn, device, reps=20): """热身 + 同步的正确计时（见 01 号文件），返回中位数（秒）。""" for _ in range(3): fn() synchronize(device) times = [] for _ in range(reps): t0 = time.perf_counter() fn() synchronize(device) times.append(time.perf_counter() - t0) return statistics.median(times) # ============================================================================= # 第 1 部分：编译一个小模块 # ============================================================================= class Block(torch.nn.Module): """ y = x + W2 · GeLU(W1 · RMSNorm(x))：Transformer 前馈子层的一个简化版本。 其中矩阵乘法（W1、W2）本来就由高度优化的库完成；RMSNorm、GeLU、残差相加这些 逐元素 / 归约运算，才是编译器可以融合、减少显存读写的地方。 """ def __init__(self, d): super().__init__() self.g = torch.nn.Parameter(torch.ones(d)) self.w1 = torch.nn.Linear(d, 4 * d, bias=False) self.w2 = torch.nn.Linear(4 * d, d, bias=False) def forward(self, x): h = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5) * self.g h = self.w1(h) h = 0.5 * h * (1 + torch.tanh(0.7978845608028654 * (h + 0.044715 * h ** 3))) return x + self.w2(h) def demo_compile_module(device): """演示 1：编译前后的速度与结果。""" print_header(f"演示 1：编译一个小模块（设备：{device}）") torch.manual_seed(0) d = 1024 model = Block(d).to(device) x = torch.randn(16, 512, d, device=device) try: compiled = torch.compile(model) with torch.no_grad(): t0 = time.perf_counter() out_c = compiled(x) synchronize(device) t_compile = time.perf_counter() - t0 out_e = model(x) print(f" 首次调用（包含编译）：{t_compile:.1f} 秒") print(f" 编译前后结果的最大差异：{(out_c - out_e).abs().max().item():.1e}") t_e = bench(lambda: model(x), device) t_c = bench(lambda: compiled(x), device) print(f" 普通执行：{t_e * 1e3:.2f} ms；编译后：{t_c * 1e3:.2f} ms（加速 {t_e / t_c:.2f} 倍）") except Exception as e: print(f" torch.compile 在这台电脑上不可用（{type(e).__name__}），跳过。") return print(""" 矩阵乘法本来就快，编译器省下的主要是 RMSNorm、GeLU、残差相加等"内存受限"运算的读写。 加速比取决于这些运算在总时间中占多少：矩阵越大，矩阵乘法占比越高，编译带来的相对收益越小。 编译有一次性的代价（首次调用很慢），所以它适合"反复运行同一段代码"的场景，比如训练循环。""") # ============================================================================= # 第 2 部分：图中断 # ============================================================================= def fn_with_break(x): """中间有一次 .item()：需要把 GPU 上的值拿到 Python 里，Dynamo 只能在这里把图切断。""" y = torch.sin(x) * 2 if y.sum().item() > 0: # 依赖数据的 Python 控制流 → 图中断 y = y + 1 return torch.cos(y) def fn_without_break(x): """同样的逻辑用张量运算表达（torch.where），整个函数是一张图。""" y = torch.sin(x) * 2 y = torch.where(y.sum() > 0, y + 1, y) return torch.cos(y) def demo_graph_breaks(): """演示 2：图中断。""" print_header("演示 2：图中断（graph break）") x = torch.randn(1000) for fn in [fn_with_break, fn_without_break]: dynamo.reset() ex = dynamo.explain(fn)(x) print(f" {fn.__name__:<18}：{ex.graph_count} 张图，{ex.graph_break_count} 次图中断") print(""" Dynamo 追踪的是"张量运算"。遇到它无法放进图里的东西 —— 依赖张量数值的 Python if、.item()、 print、调用不支持的第三方库 —— 就只能把图切断：前半截编译、回到 Python 执行、后半截再编译。 每次中断都会打断融合，还可能引入 GPU→CPU 的同步。 排查方法：torch._dynamo.explain(函数)(输入)，或者设置环境变量 TORCH_LOGS=graph_breaks。 （torch.compile(..., fullgraph=True) 会在出现图中断时直接报错，适合用来确保整个模型被编译成一张图。）""") # ============================================================================= # 第 3 部分：重新编译 # ============================================================================= def demo_recompile(): """演示 3：输入形状变化时的重新编译。""" print_header("演示 3：形状变化与重新编译") dynamo.reset() counters.clear() f = torch.compile(lambda x: torch.sin(x * 2) + 1) for n in [10, 20, 30, 40, 50]: f(torch.randn(n)) print(f" 用 5 种不同长度的输入调用：一共编译了 {counters['stats']['unique_graphs']} 张图") print(""" 第一次按固定形状编译；第二次遇到新形状时，编译器意识到这个维度会变，就重新编译一个"动态形状"的版本， 之后的新长度都复用它。所以通常只会多编译一次。 但如果每次调用的形状都不同、又触发了其他需要特化的条件，就可能反复编译、越跑越慢。 训练时常见的做法：把序列填充 / 截断到固定长度，让形状保持不变。""") def demo_limits(): """演示 4：编译器做不到什么（文字总结）。""" print_header("演示 4：编译器的边界") print(""" torch.compile 擅长的： - 把相邻的逐元素运算、归约运算融合成一个核函数（本文件演示 1，第 5 讲 03 号文件）； - 把逐元素运算"挂"在矩阵乘法的前后（有时能融合进矩阵乘法的收尾阶段）； - 减少 Python 开销、自动选择核函数的配置（mode="max-autotune" 会实测多种分块大小）。 它通常做不到的： - 需要**改变算法**的优化。标准注意力 → FlashAttention 不是"把运算合并"，而是用在线 softmax 重新组织计算顺序（第 5 讲），编译器不会自己发明这种算法。 （PyTorch 会识别出"注意力"这个模式，替换成现成的融合核函数，比如 F.scaled_dot_product_attention。） - 跨越多个 GPU 的通信与计算重叠、MoE 的路由等系统层面的优化。 所以大模型的高性能代码通常是"编译器 + 少数手写核函数"的组合： 大部分交给 torch.compile，最关键的热点（注意力、归一化、MoE 分发、量化矩阵乘法）用 Triton / CUDA 手写。""") def main(): device = get_device() demo_compile_module(device) demo_graph_breaks() demo_recompile() demo_limits() print("\n全部演示运行完毕。") if __name__ == "__main__": main()