最近在试着用 PyTorch 2.0 的 torch.compile 优化一个 7B 参数的 LLaMA 微调脚本，想着能提速顺便省点显存。结果第一次编译就把 24G 显存干爆了，之前用原生 PyTorch 还能勉强跑个 batch size=1……查了文档说可以用 mode=“reduce-overhead” 或者 “max-autotune”，但具体区别还是云里雾里。而且我试了把 dynamic=True 加上，编译时间反而更长了，不知道是不是理解错了。有没有大佬实际在生产里用过 torch.compile 的？对于大模型微调，到底怎么设置才能既提速又不让显存炸？还有那个 graph break 的 warning 影响大吗？真心求个实操经验。
楼主
2026-07-18
PyTorch 2.0 编译大模型时显存爆了，torch.compile 到底怎么用才省显存？
请 登录 后发表回复
全部回复
共 163 条
2楼
2026-09-22
我上次也踩了类似的坑，torch.compile 第一次跑会额外占不少显存做图捕获和缓存，24G 跑7B确实悬。reduce-overhead 主要省的是 kernel launch 开销，max-autotune 会疯狂试不同 kernel，编译期显存和时间的代价都更大，省显存反而不明显。dynamic=True 对变长输入友好，但会触发多次重编译，编译时间拉长很正常。真要省显存，先试试 mode="reduce-overhead" 配合 gradient checkpointing，再把 batch 压到1、seqlen 降下来，graph break 太多的话 compile 基本白搭。
3楼
2026-09-25
我上个月也在7B模型上折腾过torch.compile，显存炸的原因大概率是编译时把整个计算图展开后产生的中间激活和缓存，特别是max-autotune会尝试很多kernel变体，显存开销比eager模式高不少。reduce-overhead主要是用CUDA graphs减少kernel launch开销，对显存帮助有限，真正省显存还得靠gradient checkpointing配合compile，或者把mode设成default再手动调dynamic。dynamic=True编译变慢是因为它要对不同shape做多次编译，如果你输入长度固定就别开，反而拖效率。生产里我一般先用torch.compile(mode="default")跑通，再逐步试reduce-overhead，max-autotune只在推理且显存充裕时用。另外可以试试fullgraph=False让它在不支持的地方回退，不然一个算子失败整个图就白编了。你那个graph break的问题，可以设torch._dynamo.config.suppress_errors=True先看日志，定位到具体哪行代码导致断图再针对性改。
4楼
2026-09-26
reduce-overhead 主要是靠 CUDA Graphs 减少 kernel launch 开销，但显存反而会多吃一块，max-autotune 更激进，编译时各种调优，爆显存很正常。7B 微调我一般不开 compile，先上 gradient checkpointing + 8bit optimizer，显存稳住再说提速。dynamic=True 会让它走动态 shape 分支，编译时间变长是预期内的，除非你序列长度变化特别大，不然别开。graph break 太多的话 compile 基本白搭，建议先用 TORCH_LOGS=graph_breaks 看看哪里断了。