最近把之前写的一个图像分类项目从PyTorch 1.13迁移到2.0,听说torch.compile能白嫖加速,就试着给模型加了个@torch.compile。结果发现训练速度反而慢了20%,报错还一堆,比如什么“dynamic shape not supported”。我用的就是标准ResNet50,输入尺寸固定,数据加载也没啥花活。是不是我姿势不对?还是说这玩意儿只对某些特定场景有效?求有经验的老哥指点一下,到底该怎么用compile才能不翻车?
PyTorch 2.0的torch.compile到底能不能直接加速老项目?踩坑了
全部回复
共 152 条说实话你这个问题我上个月刚踩过一模一样的坑,ResNet50加compile反而慢20%太正常了,我连输入尺寸写死都试过,照样报dynamic shape的错。后来翻了下源码和issue区,发现torch.compile对静态shape的检测比想象中敏感,你DataLoader里如果有个随机resize或者归一化用了不同shape的tensor,它就会频繁recompile,那开销比省下的算力还大。我最后解决的办法是把compile放到推理阶段,训练保持eager模式,推理时固定batch和分辨率,再用mode="reduce-overhead"和fullgraph=True,速度确实能提个15%左右。但如果你训练代码里用了自定义loss或者比较复杂的控制流,建议别硬上compile,它现在对动态图的支持真的还很鸡肋。另外老项目迁移时最好先把torch版本升好,单独跑通forward和backward,再一步步加compile,别一上来就全链路套,不然报错信息能把人逼疯。总之这功能不是银弹,除非你模型结构特别规整、训练循环极其简单,否则前期调优成本可能比收益还高。
torch.compile对ResNet这种静态图确实提升有限,我之前试过在EfficientNet上也就快10%左右,你慢了20%大概率是CUDA graph和显存分配的开销没摊薄。建议先试试torch.compile(model, mode="reduce-overhead"),然后确保整个训练循环里没有Python原生list或dict操作,尤其dataloader返回的tensor维度别有任何隐式变化。另外报错dynamic shape很可能是你用了mixup或者cutout这类在线增强,导致batch内shape不统一,建议先关掉验证一下。
我试过类似的情况,torch.compile对固定shape的CNN其实提升有限,尤其ResNet这种老结构,计算密度不低,编译开销反而盖过了收益。你慢20%大概率是图编译和算子融合的启动开销摊不平,建议先把mode设为reduce-overhead,或者只compile model的前向部分,别包整个train_step。另外报错dynamic shape大概率是某些op内部有隐式view或reshape,你检查下有没有用Python int做索引的地方,换成tensor就好。老实说这功能更适合transformer或者大模型那种重计算场景,老项目迁移真不如先试试AMP混合精度,见效快还省事。
这情况我太熟了,torch.compile不是无脑加的,默认模式对动态shape特别敏感,你虽然觉得输入固定,但loss或中间层如果有非张量操作触发重编译,反而更慢。建议先开torch.compile(model, mode="reduce-overhead")试试,或者用torch._dynamo.config.suppress_errors=True先跑通,再慢慢排查是哪个算子导致的回退。另外训练阶段收益本来就小于推理,你这20%的倒退多半是编译开销没摊薄,老项目还是先确认算子兼容性再上吧。
小模型和短训练时长的场景真的别硬上,compile启动开销和显存占用都够喝一壶的,我上次也这样翻车了。
小模型和短训练步数真没必要上compile,开销都在图优化和编译上了,跑长训练或者大模型才能回本。
说实话你这情况我太熟了,之前我拿torch.compile跑一个分割模型也这样,ResNet这种结构其实对编译优化不太敏感,它本来就已经很成熟了,CUDNN后端早就把卷积路径调得差不多了,你编译个半天等于白忙活。而且2.0刚出那会儿编译的overhead特别高,尤其是第一次跑要花大量时间做graph capture和triton kernel生成,如果你训练步数不够多,这点时间根本摊不回来。你那报错dynamic shape,多半是DataLoader里有个别batch的size跟预设的差一点点,比如最后一批drop_last没设,或者label里有什么变长操作,torch.compile对shape的严格程度超出你想象。我现在用下来感觉这玩意最爽的场景是那种带复杂控制流或者自定义算子的模型,比如NLP里那种分支多的,还有推理时的batch=1,能让它把图整个融合起来。你要真想让它加速,建议先关掉cudnn的benchmark,把batch size调成8的倍数,还有设置mode="reduce-overhead",然后跑个几十步看下warmup之后的速度,别一上来就全量训练。反正我的经验是,老项目如果本来就是标准结构,直接上compile不如把精力花在调数据加载和混合精度上,那才是实打实的收益。
老项目直接套compile确实容易翻车,我这边模型带动态shape也是负优化,建议先跑一遍静态shape的基准测试再决定。
这情况太真实了,我试过在分割模型上开compile,显存直接爆了,后来发现得先跑几步warmup让CUDA图缓存起来才行。你ResNet50这种静态输入按理说最友好,试试把torch._dynamo.config.suppress_errors=True加上,然后设个torch.set_float32_matmul_precision('high'),有时候是跟AMP混用才有效果。另外老项目里有些自定义op或list/tuple操作容易触发graph break,一break就白搭了,建议先看下编译日志里有没有warning。
说实话你这情况我太熟了,刚升2.0那会儿我也这么干过,ResNet50加compile结果原地踏步甚至倒退,后来查了半天才发现是CUDA graph和我的自定义loss里某些操作冲突了,导致每次反向都重新编译。你这“dynamic shape not supported”大概率不是输入尺寸的问题,而是模型里某个op的输出shape在运行时发生了微妙变化,比如BatchNorm在eval和train切换时的行为,或者某些reshape操作依赖了tensor的实际值。建议你先用torch.compile的mode="reduce-overhead"试试,同时把model.eval()和model.train()分开测,别一上来就默认全流程加速。另外有个坑是,compile对显存占用不友好,如果你batch size本来就卡在显存边缘,编译后的临时缓存可能直接让你OOM,那速度反而更难看。我现在的做法是,先跑一遍纯eager模式记录下每个op的时间,再用torch.profiler看compile模式下哪些kernel被替换了,如果发现替换后的kernel没比原来快多少,那这层模型就不值得compile。还有个小技巧,如果你的数据加载用了num_workers,记得把compile放在模型定义之后、DataLoader之前,不然有时候会莫名触发多次编译。总之这玩意儿不是白嫖,更像是你得先给它铺好路,把动态性都消除干净才有效,尤其老项目里那些花里胡哨的写法,可能比你想的更不兼容。
踩坑+1,小模型和动态shape真别硬上,compile对大模型和静态图收益才明显。
老项目先别全局开,挑计算密集的模块局部试,配合torch.compile模式里的reduce-overhead可能好点。
torch.compile对ResNet50这种CNN确实收益不大,它主要吃香的是Transformer和动态图结构,我试过在分割模型上也是负优化。你检查下是不是CUDA graph没生效,或者batch size太小导致编译开销摊不平,建议先把torch._dynamo的日志打开看看graph break在哪。另外试试mode="max-autotune"或者关掉dynamic=True,有时候默认设置反而触发不必要的shape检查。老项目想白嫖不如先升级到2.x然后开channels_last和AMP,这个提升更稳。
我跟你情况差不多,也是ResNet50,刚开始compile直接给我整不会了,后来发现问题出在没关cudnn的benchmark模式,那个跟torch.compile的graph优化有冲突,关掉之后速度就上来了。另外你那个dynamic shape的报错,八成是模型里有个别层输出尺寸在跑的时候会有细微变化,比如最后的global average pooling之后加了个view或者reshape,建议把输入tensor用torch.zeros跑一遍trace,把所有中间shape都打出来看看。还有一个坑是compile默认的mode是default,它对内存开销比较敏感,小batch下反而容易慢,你试试mode="max-autotune"或者"reduce-overhead",后者显存够的话能明显提速。不过说实话,对于这种标准CNN,如果batch size不大(比如32以下),compile的提升真的很有限,有时候还不如直接开amp混合精度来得实在。我现在是这么用的:训练阶段干脆不compile,只把推理部分的模型compile一下,配合torch.inference_mode,这样稳定性和速度都兼顾了,你可以试试这个思路。
说实话你这个情况我太熟了,刚升2.0那会儿我也是直接无脑套compile,结果跟你一样,训练直接卡成ppt,后来才发现这玩意儿默认开的是reduce-overhead模式,对某些卡和CUDA版本反而有额外开销。你试试torch.compile(model, mode="max-autotune"),然后配合torch._dynamo.config.suppress_errors=True,至少先把报错压下去,但速度提升真不一定有。另外ResNet50这种CNN其实在1.x版本里已经被cudnn benchmark优化得很透了,compile主要强在融合算子,对transformer或者有大矩阵运算的模型收益更明显,你拿它去编CNN可能真是反向优化。还有个坑是batch size如果不大,比如32以下,compile的图优化启动时间占比太高,反而拖慢总时长,可以试试把batch调大点或者用accumulation steps模拟大batch。最后建议你直接看下torch.compile生成的inductor日志,如果显示很多“skipped”或者“unfused”,说明你的网络结构根本没被有效改造,那不如干脆关掉,省心。
我之前也踩过这坑,后来发现torch.compile默认会做很多动态shape的假设检查,你固定尺寸反而容易触发它过度优化,试试加mode="reduce-overhead"或者把dynamic=True参数显式关掉,能省不少事。另外老项目里如果有自定义loss或前向逻辑里用了Python控制流,compile基本就白搭,甚至会拖慢,建议先用torch.profiler定位下到底瓶颈在哪,说不定是CUDA graph和你的数据加载不匹配。还有个小技巧,compile前先把模型eval和train模式下的输入都跑一遍warmup,能减少很多诡异的报错。
先关掉动态shape再试,小batch下compile反而慢,把batch调大点或者用模式reduce-overhead试试。
torch.compile第一波跑的时候会做graph capture和算子融合,这个编译开销有时候比训练本身还大,你那个慢20%可能就是因为这个,多跑几个epoch再看看曲线,通常后面会追回来。另外ResNet50这种CNN其实不太吃compile的收益,它主要利好Transformer或者带动态控制流的模型,你试试把batch size调大点,或者用mode="max-autotune"看看。报dynamic shape的话,检查下有没有把tensor转list或者用到Python int做索引,固定尺寸就尽量全用torch tensor操作。
我之前也踩过同样的坑,后来发现torch.compile对动态shape特别敏感,你如果数据加载和模型输入都是固定的话,可以试试把dynamic=False显式传进去,或者关掉cudagraphs看看。另外老项目里有些自定义op或者第三方库没适配2.0的编译,也会导致回退到eager模式,速度反而更慢。建议你先用torch.compile的mode='reduce-overhead'跑个benchmark对比,如果还是负优化,可能你项目里某些操作(比如F.interpolate)触发graph break太频繁,这时候不如只对backbone部分做compile。我自己的经验是,它更适合大模型和GPU资源充足的情况,小数据集小模型收益确实不明显。
我也是从1.13升上来的,第一次跑torch.compile跟你一模一样的感受,慢20%算少的,我这边直接OOM了。后来翻了下源码和issue,发现这玩意儿默认开了一堆graph break和CUDA graph的优化,但老项目里但凡有个动态shape或者自定义loss,它就得频繁回退到eager模式,来回切换的开销比不编译还大。建议你先别直接@torch.compile整个模型,试试只包住backbone那部分,或者把输入张量用torch.empty显式声明一下max_shape,至少能少一半的报错。另外2.0的compile对A100这类新卡效果明显,如果是老卡比如V100,收益真的有限,甚至可能负优化。我后来是把torch.compile换成torch._dynamo.disable()包住几个不友好的模块,剩下部分再编译,速度才勉强提了10%左右。反正别指望无脑白嫖,得配合torch._dynamo.config里的dynamic参数调一调,比如dynamic=False强制静态shape,但你要是数据加载里有随机resize就得小心了。你要是代码不大,可以试试把数据管线改成纯tensor操作,喂给模型前固定好尺寸,应该能有改善。
说实话你这个问题我太有共鸣了,当时我拿torch.compile去跑一个老的分割模型也是这鬼样子,慢倒其次,那堆报错直接给我整不会了。后来我翻了不少issue才搞明白,这玩意儿对静态shape和连续显存布局特别敏感,你虽然固定了输入尺寸,但如果你代码里有用到Python原生list做数据拼接、或者在forward里动态改张量结构,它就会疯狂recompile,反而比eager模式还慢。我建议你先用torch.compile(model, mode="reduce-overhead")配合torch._dynamo.config.suppress_errors试试,把报错压下去看真实速度,同时检查一下你的collate_fn和dataset是不是每次都产生新shape的tensor,哪怕看起来一样,dtype或stride不同也会触发guard miss。还有个坑是别把整个train_step都塞进去,最好只compile模型本身,损失函数和优化器保持原样,我这么改完才勉强跟原来打平。另外如果你用混合精度,记得先amp再compile,顺序反了会有一堆算子融合问题。说实话,老项目真想白嫖加速,不如先试试torch.backends.cudnn.benchmark=True和通道last格式,那个提升才叫无痛。