前阵子我把一个跑了大半年的视觉推理服务从PyTorch的默认Eager模式切到了torch.compile没动一行推理逻辑端到端延迟直接掉了将近一半。这件事让我意识到很多人在PyTorch上写模型、训模型、部署模型却一直停留在“能用就行”的阶段完全没有吃到框架自带的性能红利。这篇文章就把我这次改造的完整思路写下来包括Eager模式为什么慢、torch.compile到底做了什么、训练和推理场景分别怎么用、以及我在实际项目中踩过哪些坑。适合已经能跑通PyTorch训练代码、但还没系统性做过性能优化的开发者阅读。1. 先把话说清楚PyTorch的Eager模式到底慢在哪1.1 逐个算子启动的开销比想象中大得多PyTorch默认的执行方式是Eager模式。你写一行output self.conv(x)框架就把这一个卷积算子独立地提交到GPU上执行。每一个算子的执行过程至少要经历Python解释器解析这行代码、PyTorch分发到对应的C实现、分配或获取内存、把kernel提交到GPU流、等待或异步返回。等你循环几百个算子走完一个forward真正花在GPU kernel上的时间反而只占了一部分剩下全浪费在调度和内存管理上。这个问题在模型结构很深、算子又小又碎的时候尤其明显。我举个例子一个标准的Transformer Block里面包含矩阵乘、LayerNorm、激活函数、残差连接、Dropout每一个都是独立kernel。假设一层有30个算子12层就是360个算子。Eager模式就老老实实启动360次kernel哪怕其中很多算子处理的数据量并不大启动开销占比依然很高。1.2 动态图和静态图的路线之争PyTorch一直坚持动态图路线好处是写代码特别自由你可以用if控制流、用Python循环、任意Dict和List结构调试的时候直接print中间张量非常符合人的直觉。但代价就是计算图是边跑边建立的框架很难提前看到整张图也就没法做跨算子的融合优化。相反静态图方案经典的图编译推理引擎、ONNX导出后的优化图因为提前拿到了整张计算图可以做算子融合、内存复用、布局优化执行效率比单个kernel一个个跑高很多。PyTorch团队后来显然也想通了动态图的开发体验不能丢但静态图的性能红利也得拿。于是就有了Dynamo和TorchInductor这套编译栈也就是torch.compile的底座。理解这一点是全文的基础torch.compile本质上就是把你的Python代码在运行前“瞥一眼”尝试把能抓到的算子拼成一张静态图再交给后端生成优化过的kernel。抓不到的部分就叫“图断裂”会回退到Eager执行。2. torch.compile的几种运行模式与后端以及我为什么默认选reduce-overhead2.1 三种模式和它们的侧重点torch.compile接受一个mode参数我常用的是这三种差别非常明显mode核心手段典型收益代价defaultTorchInductor生成融合过的Triton Kernel访存型算子融合明显计算密集型也有提升编译时间较长reduce-overhead在default基础上额外使用CUDA Graph捕获大幅减少kernel启动开销小算子密集的模型提升巨大显存占用略增部分动态Shape不友好max-autotune对每个算子做更激进的自动调优可能达到最优kernel配置编译时间极长显存占用高我第一次上手就直接试了max-autotune结果一个中型模型编译了快20分钟而且自动调优期间显存吃满差点把同机器上的其他任务挤爆。后来学乖了先用default跑通再用reduce-overhead看收益实在需要压极限才上max-autotune并且只对推理服务用。2.2 后端Inductor和CUDA Graph的关系很多人把reduce-overhead理解为“自动变大batch”其实是误解。它做的事情是把整张计算图的所有kernel一次性捕获进一个CUDA Graph之后每次执行只提交一次launch。GPU不需要CPU反复下发几千个kernel直接按图执行这个开销降得非常明显。对于算子小而密的模型比如小型CNN或轻量Transformer收益能到20%到40%和你batch大小没有直接关系。不过CUDA Graph有一个硬性要求图被捕获后里面每个tensor的内存地址和shape基本固定。一旦你的输入动态变化比如batch忽大忽小、序列长度每句不同图就得重新捕获这个重捕过程不仅慢还会导致显存中留下很多旧的内存碎片。TorchInductor是当前PyTorch默认代码生成后端负责把FX Graph转成针对GPU的Triton Kernel。它做了很多经典优化把elementwise算子融合成一个kernel、把相邻的residual和normalization合并、按kernel的访存模式重新安排执行顺序。我实测下来对LayerNorm、SiLU、残差相加这类访存密集型算子Inductor融合之后往往有2到3倍的单算子性能提升。3. 把一段训练代码改造成torch.compile从包装模型到处理动态Shape3.1 最基本的接入方式接入编译模式比很多人想象中简单核心就是一行import torch model build_my_model() model torch.compile(model, dynamicFalse, modereduce-overhead)然后该咋训练就咋训练。optimizer不需要动DDP也不需要额外配置PyTorch会自己在编译后的forward外面包好梯度相关逻辑。我第一次实验的时候心里其实很虚总觉得加了这一行代码会发生什么玄学问题。事实证明大部分标准模型结构都能直接跑前提是你写的forward里面没有太多“图抓不住”的操作。以PyTorch 2.x版本为例一个典型训练循环改造后长这样import torch import torch.nn as nn from torch.cuda.amp import autocast, GradScaler model nn.Sequential( nn.Linear(768, 3072), nn.GELU(), nn.Linear(3072, 768), nn.Dropout(0.1), ).cuda() model torch.compile(model, dynamicFalse, modereduce-overhead) optimizer torch.optim.AdamW(model.parameters(), lr3e-4) scaler GradScaler() criterion nn.CrossEntropyLoss() def train_step(batch_x, batch_y): optimizer.zero_grad() with autocast(dtypetorch.bfloat16): logits model(batch_x) loss criterion(logits, batch_y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() return loss.item()需要注意一点reduce-overhead配合autocast时GradScaler照常使用。CUDA Graph捕获过程中不能和CPU频繁同步所以不要在任何被编译的模型内部写loss.item()或torch.cuda.synchronize()。这类同步操作一旦出现在图中间Dynamo大概率会把它变成一个图断裂点编译后的收益直接打折。3.2 动态Shape是常见的隐形成本来源如果你的训练数据里sequence length是固定的那dynamicFalse是最省心的选择CUDA Graph捕获一次就能一直复用。但NLP场景里padding之后每个batch的序列长度可能还是不一样这时就进入动态Shape的范畴。我建议第一步不是立刻打开dynamicTrue而是先检查你的DataLoader是不是真的在制造不同的shape。很多情况下只要把padding策略改成“按batch内最大长度动态padding”整个训练流程里实际shape变化的数量也就几种。这时候完全可以保留dynamicFalse或者用torch._dynamo.mark_dynamic把某一个维度标记为动态让Dynamo对那个维度做专门处理。model torch.compile(model, dynamicFalse, modereduce-overhead) torch._dynamo.mark_dynamic(batch_x, 1) # 标记序列长度为动态维度这里dynamic参数名在不同PyTorch小版本上略有差异2.2以后更推荐直接在torch.compile里指定dynamicTrue或者在导出时用专门的动态维度API。核心理解就一条动态Shape一旦出现编译缓存可能频繁失效CUDA Graph会反复重捕性能反而可能不如Eager。4. 推理服务的实战改造固定Shape换来的延迟稳定性4.1 动态batch导致的不只是变慢还有抖动我在生产环境里维护过一个人脸关键点检测服务输入是摄像头截取的单张图早期代码为了灵活直接把原图resize到正方形再进模型。每个请求的图片分辨率都不同模型内部的计算shape就一直在变。用Eager模式跑倒是看不出太大问题毕竟每次都是新kernel。但切到reduce-overhead之后问题立刻暴露每个新shape都会触发CUDA Graph重新捕获于是服务延迟忽高忽低P95反而比Eager更差。后来我把服务改成固定推理尺寸所有请求统一resize到256x256的tensor模型编译保持dynamicFalse、modereduce-overhead。改动上线后单图推理延迟的P95从25毫秒左右降到13毫秒左右而且波动小了很多。这个优化思路说起来非常朴素就是“让编译图永远复用同一套内存布局和kernel配置”但带来的收益却非常可观。对推理服务来说稳定P95往往比降低P50更重要。固定输入Shape配合CUDA Graph最大的价值不是极限性能而是把延迟方差压下来。线上服务追求的是“每个请求都快且差不多快”而不是“大多数请求飞快、偶发请求卡顿”。4.2 推理时记得用的两个东西inference_mode和编译预热我在评测很多同学的推理代码时发现他们仍然在用torch.no_grad()甚至部分代码连这个都没加。torch.no_grad()只是不追踪梯度但模型内部的自动混合精度、算子选择、内存分配仍然按照训练模式来处理。正确的做法是在推理阶段用torch.inference_mode()它能更彻底地关闭梯度相关的bookkeeping和torch.compile配合起来效果更好。另一个被忽略的小点是“编译预热”。torch.compile在第一次前向时要做图捕获、Triton编译、CUDA Graph捕获第一次延迟会明显偏长。你不希望线上第一个真实请求去背这个黑锅。所以模型初始化之后我会用一批随机的假数据先跑十几遍forward确保编译缓存真正落定再打开服务对外接收流量。model.eval() compiled_model torch.compile(model, dynamicFalse, modereduce-overhead) # 预热多跑几轮让CUDA Graph捕获和kernel autotune完成 with torch.inference_mode(): dummy torch.randn(1, 3, 256, 256, devicecuda) for _ in range(20): _ compiled_model(dummy)如果把编译后的模型继续用于训练比如fine-tune预热时一定不要用torch.inference_mode()否则模型内部一些训练态逻辑如Dropout的随机状态不会被正确建立可能导致后续训练时结果对不上。我自己的做法是训练代码里不单独预热直接让前几个step吃点编译开销推理服务里才专门预热。5. 我踩过的坑和判断标准什么情况下不要用torch.compile5.1 图断裂最常见也最容易忽略Dynamo是一个Python字节码层面的图捕获器它会把你的forward函数从头到尾扫描一遍遇到它不能理解的操作就在那里“断”一下断点之后继续用Eager模式执行。图断裂本身不会让程序出错但次数一多编译图的覆盖率就下降性能自然上不去。我见过很多自定义模型里藏着大量“不可编译”的写法典型的有直接依赖Python对象身份判断的ifforward里动态创建新的层结构比如循环里append一个新nn.Module调用外部Python库处理张量比如用纯Python写注意力掩码、排序、索引在forward内部做tensor.item() / tensor.tolist()之类的CPU同步操作判断一段模型到底编译了多少、断了多少我一般用这个命令from torch._dynamo import explain explanation explain(model, dummy_input) print(explanation)输出结果里会给出graph breaks的位置、原因、以及断成几段。我改造那个视觉服务的时候最初用explain发现一个模型居然断成7段后来把里面一个索引计算从Python纯逻辑换成torch.gather段数立刻降到了1段性能提升非常明显。5.2 自定义算子与CUDA Graph的组合问题如果你的模型里包含了自定义CUDA扩展、第三方算子库或一些老旧的ONNX算子映射直接上reduce-overhead很容易踩雷。CUDA Graph捕获要求所有kernel都被正确记录而第三方库的kernel不一定记录在标准流里捕获阶段可能直接报错或者捕获到一张残缺的图。这种情况下我建议先用modedefault只做Inductor融合别碰CUDA Graph或者把这些自定义算子单独放到一个不被编译的子模块里。PyTorch官方在编译性方面做了很多兼容处理比如遇到没法处理的算子会回退但回退依然意味着性能损失。我的选择标准很简单自定义算子在整体计算中占比很低就直接接受图断裂占比很高就优先考虑用原生算子重写这部分而不是去绕开编译。5.3 什么时候干脆别用torch.compile不是所有模型都适合编译。PyTorch本身迭代速度很快但torch.compile的适用边界还是很明显的下面这些场景我建议先别上或者至少先做个小规模验证模型很小、batch又小整体都在微秒级执行torch.compile的模块初始化开销和显存占用可能比省下来的计算还多。模型包含大量动态控制流、稀疏计算、动态索引图捕获的性价比太低。典型如某些推荐系统的Embedding聚合逻辑。你已经用了厂商专用推理引擎比如导出的图优化推理格式并且性能已经足够不需要再用torch.compile再叠一层。两者在算子布局上的假设可能不一致混合使用反而会有额外内存拷贝。做研究实验时频繁修改模型结构每次修改都会触发重新编译等待编译的时间会明显拖慢迭代节奏。这时候Eager模式的灵活性才是真正的刚需。我现在的习惯是技术验证阶段先用Eager快速跑通等模型结构固定、准备做规模化训练或部署时再引入torch.compile。并且在每次引入前都先记录Eager模式的基线指标比如训练吞吐、推理延迟、显存占用编译后对照同一套数据做A/B不让感觉代替数据。最后再分享一个细节。很多人不知道torch.compile的编译产物是可以缓存的默认的缓存目录在~/.cache/torch/compile。如果你重启进程后还要重新编译检查一下这个目录的写入权限以及环境变量TORCHINDUCTOR_CACHE_DIR是否被设置成了某个不可写的路径。我在服务器上曾经遇到过容器重启后每次都要重新编译的情况排查了半天就是挂载卷把缓存目录覆盖了。这个坑不大但相当费时间。总之torch.compile是PyTorch这两代版本里最值得投入时间研究的性能工具但它的收益建立在理解图、理解内存、理解执行模式的基础上。先把Eager模式为什么慢搞明白再去调torch.compile你会省掉很多无效尝试。