PyTorch torch.compile 排障实战手册:从入门上手到深度调优的完整指南
本文以 PyTorch 官方仓库中的技术文档 torch.compile, the missing manual 为主体骨架整理扩充而成。该文档聚焦一个核心问题:当你在自己的模型上启用
torch.compile后遇到了性能不及预期、编译报错、内存超支、分布式超时或输出错乱时,应当如何系统化地定位与解决。文中所有配置项、命令与结论均对照当前仓库源码(torch/_dynamo/config.py、torch/_inductor/config.py 等)逐一验证,并在原文基础上补充了源码级注释与最新演进说明。
torch.compile 是 PyTorch 2.x 引入的即时编译器(PT2 / Dynamo + Inductor),目标是让既有的 eager 模型在几乎不改代码的情况下获得加速。但它是一套复杂且较新的软件,你在使用中大概率会遇到"成长的烦恼"。本手册面向对 PyTorch 内部机制了解不多、但熟悉自己模型、并愿意与 PyTorch 开发者通过 issue 协作的技术型用户。你不必从头到尾通读,可按需跳转:模型启不启得动看"预期管理";跑起来不对看"编译器崩溃"与"输出是垃圾数据";编译太久看"编译时间过长";结果不理想看"性能不达预期";分布式卡死看"NCCL 超时"。
版本适用说明:原文 NOTICE(Jul 5, 2024)说明其针对当时最新的 PyTorch nightly。本文所引用的仓库代码已明显晚于该时间点,例如
inline_inbuilt_nn_modules已在 torch/_dynamo/config.py 中变为默认开启且被标记为 deprecated("does not do anything, inline_inbuilt_nn_modules is always True"),optimize_ddp也扩展出了更细粒度取值。因此文中凡涉及默认值,均以当前仓库代码为准,并尽量同时说明演进背景。
三种落地难度等级(The three regimes of enablement)
PyTorch 中模型架构千差万别。虽然 torch.compile 定位为通用编译器,但实际帮助用户接入时,模型通常落入以下三个等级之一:
- It just works(开箱即用):通常出现在三种情况:(1) 模型是社区重点打磨过的;(2) 模型不复杂、写法是简单惯用的 IDIOMATIC 代码;(3) 模型从设计之初就以 torch.compile 友好为目标编写(如 gpt-fast、torchao 一类的仓库)。
- It works with a little work(略作改造即可用):真实业务里常有各种"毛刺",例如第三方库使用了花哨的 Python 特性,或标准 Transformer 中有一两处建模创新。这类模型本质上与能跑好的模型相似,你只需少量改写代码绕开个别编译器 bug,即可用较小投入接入 torch.compile。
- It's going to be a slog(会是一场持久战):你的模型在做一些"奇怪"的事——例如高度优化的 eager 代码在 backward hook 里做文章,或利用数据相关计算做稀疏优化,或代码充分利用 eager 特性在 Tensor 与纯 Python 计算之间来回穿梭。这类模型通常需要与 PyTorch 团队合力修复数十个 bug 才能真正拿到性能数据,之后还可能要处理正确性与性能问题。中途放弃是完全合理的。但值得一提的是:torch.compile 的持续开发恰恰在重点改善这类场景(Meta 大量推荐模型就落在等级 3),如果你愿意投入并与其"联合开发",隧道尽头可能藏着可观的加速收益。
什么应当被期望能够编译
一个训练脚本中有多个可独立编译的部分,实际上大部分都能编译,文档给出如下细分:
- 编译 nn.Module 是传统且最成熟的主战场。被编译的模块会产生优化的 forward 与 backward,等价于你手写了一个带优化前向/反向的自定义 autograd Function。两种常见写法:
torch.compile(...)(nn_module)(构造一个包裹原模块的 wrapper),或只编译 forward:nn_module.forward = torch.compile(...)(nn_module.forward)。 - 编译优化器(compiled optimizers)也已受支持:把
torch.compile包在 optimizer step 调用外即可。优化器曾有一些特殊问题(截至 2024-07 已基本解决,Adam、AdamW、SGD、RAdam、Adamax 均预期可用)。两点关键:- 优化器常对纯 Python int/float 做大量计算,而 Dynamo 对这类标量捕获不佳。多数优化器提供 capturable 变体(把计算放到 Tensor 上;capturable 优化器对 eager 也有益,因为 CUDA graphs 依赖它)。LR scheduler 通常是 float 学习率,必须包成 Tensor 才能被 torch.compile 处理(H2 2024 有计划移除该限制,仓库 issue #107277 跟踪)。
- 理论上优化器可以逐个参数朴素更新、靠编译器做水平融合,但实践上应直接使用 foreach 内核——让 Inductor 逆向工程水平融合对编译时间很不友好(大量参数要融合 + 融合 pass 的非线性开销,双重打击)。编译优化器的主要加速来源仍是对这些更新的垂直融合。
- 编译 autograd(compiled autograd):编译模块通常顺带给了编译后的 backward;但你也可以通过
torch._dynamo.compiled_autograd(仓库实现位于 torch/_dynamo/compiled_autograd.py)直接编译backward()执行的 autograd 图。三种推荐使用场景:(1) 开箱即有收益——accumulate grad 节点可被融合进编译区域(传统 AOTAutograd 做不到);(2) forward 因动态性无法整体编译,但 backward 每轮形状恒定;(3) 用到了 hooks 这类无法提前编译的复杂 autograd 特性(编译 FSDP 等分布式 wrapper 时尤其常见)。 - 日志打印通常会引发 graph break(它们代表 PT2 IR 当前无法表达的副作用)。可用
torch._dynamo.config.reorderable_logging_functions把某些函数(如print、warnings.warn,包括打印 Tensor 的日志函数)标记为可重排到编译区域末尾。注意:这种日志可能改变性能特性(原本不必实体的 Tensor 可能为打印而实体化),且日志只会在执行结束时打印——若 buffer 被就地修改,你会看到修改后的值。该配置在 torch/_dynamo/config.py 附近定义。
当前不支持的场景(文档明确列出):
- 单步捕获(Single step capture):把 forward、backward、optimizer step 一次捕获进同一张图的假设性策略,仍有 RFC 在推进(issue #117394),截至 2024-07 尚未实现。
- 预处理(Preprocessing):预处理往往依赖大量领域专属自定义算子;PT2 原则上能捕获自定义算子,但你不一定拿到运行时加速。截至 2024-08 更新,Meta 内部对"原生 PyTorch 预处理"兴趣渐增,若你确有需求可主动联络 PyTorch 团队。
无论遇到哪一类问题,以下调试方法普遍适用。文档建议按从省力到费力的顺序依次尝试。
用 TORCH_TRACE 采集 trace,并用 tlparse 查看
TORCH_TRACE 与 tlparse 是一对产生"编译报告"的工具。采集极其简单:
TORCH_TRACE="/tmp/tracedir" python foo.py
pip install tlparse # 如未安装(Meta 内部用 `feature install tlparse`)
tlparse /tmp/tracedir
- 分布式任务同样可用(每个 rank 一份 trace)。如果你遇到一个复杂问题却无法给出独立最小复现,把
/tmp/tracedir下的 trace 日志附到 bug report 里就能极大帮助 PyTorch 开发者。 - 安全警告:trace 日志包含你的全部模型代码(但不含权重)。模型敏感时不要外传。
- 采集开销足够低,Meta 内部甚至在 torch.compile 的生产任务中默认开启它。
tlparse 报告主要面向 PyTorch 开发者,但普通用户也能从中提取有用信息(报告中自带 help 文本):
- 实际编译了哪些模型代码(通过 stack trie,尤其适合不熟悉被编译代码库的人);
- 有多少 graph break / 多少个独立编译区域(每个编译区域是独立着色块;浅绿色代表可能被 graph break 的帧;帧数过多很可疑,提示灾难性 graph break 或代码与 torch.compile 契合度差);
- 某个帧被重编译了多少次(如
[10/0] [10/1] [10/2]——频繁重编译非常可疑,即使不是根因也值得调查); - 是否发生编译错误(出错的帧会显示为
[0/1]样式); - 某个帧生成了哪些中间编译产物(很多人喜欢看
inductor_output_code_*里的 Triton 代码:能把某个 Triton kernel 与用户代码关联起来的话,专家用户可以直接看出 PT2 在生成什么); - 特定帧的相关信息(在 compilation_metrics 中)。
若你的 PT2 模型已上线,升级 PyTorch release/nightly 后可能"莫名其妙"失效——把正常与异常版本的 tlparse 并排对比(视觉 diff)是很好的入手方式。
做消融实验(ablation)定位病灶
崩溃时根因通常一目了然;但输出是"垃圾数据"时没有着手点、也难以做出最小复现。此时第一要务是消融:重跑 workload,依次关闭编译器栈的各个部件,隔离出问题组件。两种消融路线:
路线一:逐层关闭编译器栈(改 backend 参数)。推荐依次测试三个档位:
backend="eager":仍失败说明是 Dynamo 的问题;backend="aot_eager":此档失败而 eager 不失败,说明是 AOTAutograd 的问题;backend="aot_eager_decomp_partition":此档失败而 aot_eager 不失败,说明问题出在 decompositions / partitioner。
此外:使用 mode="reduce-overhead" 时尝试去掉它(排查是否为 CUDA graphs 问题);使用 dynamic=True 时尝试去掉(排查是否为动态 shape 问题)。若怀疑与某个 FX pass 有关,可手动注释或关闭对应 config——FX pass 的入口分别在 torch/_inductor/fx_passes/joint_graph.py(joint graph,即前向+反向联合图)与 torch/_inductor/fx_passes/post_grad.py(grad 之后、decomposition 之后的图)。
路线二:按模型层级关闭编译器。把 torch.compile 调用移进模型内部即可手工实现;也可用脚本程序化地按模块禁用(文档附有 gist 补丁示例,TODO 计划正式落地)。
除了逐层对比输出(如逐层比较中间结果)之外,消融最简单省力,应当最先做。
自动 repro 生成器(minifier)
文档坦诚:多数时候 minifier 不work,你仍得另寻他路(作者猜测:能被自动复现的 bug 大多易修、早被修完,剩下的是难以自动复现的硬骨头)。但它试起来很快,值得一试。操作指引:
- 编译器崩溃时设置
TORCHDYNAMO_REPRO_AFTER="dynamo"或TORCHDYNAMO_REPRO_AFTER="aot"(aot 成功概率更高,但捕获不到 AOTAutograd 自身的 bug),祈祷生成的 repro.py 确实复现了你的问题; - 准确性问题可尝试
TORCHDYNAMO_REPRO_LEVEL=4。
检查近期 feature flag 是否变化
若某个 PT2 运行此前正常、未做代码发布却突然失效,很可能是有人改了 feature flag 改变了默认行为。开源版 PyTorch 默认不接入 feature flag 服务,但预留了扩展点,可对接自己的 flag 系统。
Bisect 二分
如果"以前能跑、现在不能跑",二分定位是常用手段(Meta 内部广泛使用)。开源环境下较痛苦:需要找到旧 nightly 下载或从源码构建对应 commit。
制作一个独立复现脚本
做 reproducer 很费功夫,没时间完全可以不做——但一个独立复现对修复 bug 的价值巨大。没有 reproducer,bug report 必须信息充分到开发者能自己从零写出复现。文档按优先级给出 reproducer 列表:
- 自包含(无外部依赖)、小型(<100 LOC) 的复现脚本——最佳;
- 自包含但较大的复现——自包含本身就是巨大加分项;
- 不自包含但对依赖版本不敏感——例如
pip install transformers后再跑脚本即能复现,也还可以接受; - 不自包含且需要大量环境搭建 / Docker image——环境越难搭,越难复现与调试。注意 Docker 让"搭环境"变容易,却让"改环境 / 使用开发者偏好的开发环境"变难,并非银弹。
正交地看:单进程复现优于多进程训练复现。对 Dynamo 相关的问题(如"不支持某 Python 构造")通常无需复现——未实现什么一目了然。
根据 bug 报告编写测试用例
拿到根因假设后,常可从零写一个测试用例去"撩拨"该问题。这需要更多 PT2 内部知识(知道用什么工具触发特定代码路径)。tlparse 是写测试的最佳帮手——用它检视真实 bug 场景中生成的 IR,能启发你如何单独触发它。按子系统分类的检查清单(文档整理,非穷尽):
- Autograd:输入 Tensor 是否
requires_grad=True?是否对输出调用了backward()? - Dynamic shapes:是否设置了
dynamic=True?是否用多种 shape 多次运行测试代码? - 自定义算子:真实流程是否涉及自定义算子?能否用 Python 自定义算子 API 复刻其关键特征?
- 配置:是否完全复刻所有配置(
torch._dynamo.config与torch._inductor.config各项、以及torch.compile的 backend/mode 参数)? - Context managers:是否复刻了活跃的上下文管理器?如
torch.no_grad、自动混合精度、TorchFunctionMode/TorchDispatchMode、activation checkpointing、compiled autograd 等。 - Tensor subclasses:是否涉及 tensor 子类?
- Partitioner:若 bug 与 partitioner 决定为反向保存什么有关,可考虑写一个带自定义 backward rule 的自定义 op,从而精确控制反向保存内容。
先问:你的程序真的应该被编译吗?
高层次地讲,只有"运行时保持同一串固定 PyTorch Tensor 操作序列"的程序才值得以有意义的方式编译。这与 TorchScript 不同——TorchScript 能捕获循环、Python 列表等语言特性并在运行时执行;而在 PT2 中,这类部分交给普通 CPython 解释器执行,torch.compile 只作用于"tensor compute"。即便简单到"用动态索引访问 Python list"或"一个长度会变的 Tensor 列表"都不受 PT2 支持。若你的模型卡在这里,或许该考虑用专门面向这类代码的工具。
先问这个问题能帮你省下大量硬啃编译器的时间:用花哨 Python 特性或操作 Python 数据结构的代码更容易触发 Dynamo 崩溃,此时正解不是"硬闯崩溃",而是"重新评估你应期望编译什么"——就算硬让它编译成功,大概率每轮都在重编译,编译变得毫无意义。
文档还给出机械化的自检方法:用 LoggingMode 之类工具对模型做多轮 Tensor 操作 trace(覆盖有代表性的输入分布)。若所有 trace 相同,基本可编译(Dynamo 的职责正是以 sound 方式捕获该 trace);若 trace 各不相同,则多半不可编译。
如果代码大量使用 Python 数据结构、又非常想用 torch.compile,核心思路是把数据结构表达成纯 Tensor。典型例子:变长序列可朴素表示为变长 Tensor 的 Python list——这无法编译(变 batch 会改变 list 长度,不被允许)。但 nested tensor 这类抽象把张量打包进单个 Tensor、外加记录各序列边界的元数据,这种打包表示就是 torch.compile 友好的。旁注:固定 list 长度时确实可以编译变长 Tensor 列表,但列表很大时会变成大量 Tensor 与大量逐张量计算,编译器可能跑得很慢(参见"循环被展开")。
三类编译器崩溃
必须区分三类崩溃:
- 不完备导致(incompleteness):报错形如
torch._dynamo.exc.Unsupported: call_method,偶尔是从 Dynamo 代码里抛出的 AttributeError(细看其实是漏处理了某种情况)。这类问题极其常见,但不代表深层毛病,只是待补的功能工作。直接附上错误信息与异常即可;可用fullgraph=False允许 graph break 绕过,或用torch.compiler.disable(仓库实现见 torch/_dynamo/decorators.py)选择性禁用模型部分编译。 - 真 bug(bugs):确实出了岔子,bug report 中信息/复现越多越好。
- 用户侧问题(user problems):错误信息往往很长很全,需要一定的用户教育。例如处理 data-dependent 算子时会遇到 GuardOnDataDependentSymNode 类错误——你试图对某个由数据推导出的符号量建立 guard 时就会触发,有时需要框架修复,有时需要换一种代码写法。
graph break 数量会实质影响你见到的崩溃形态。频繁 graph break 的模型会产生大量小块动态代码片段,与一次全捕获的单一模型相比,对编译器的压力点完全不同(插入 graph break 与重建 Python 状态的逻辑更复杂,更容易踩到 bug)。尽可能减少 graph break 能降低崩溃概率。
运行时崩溃也可能源于生成的代码有缺陷。TORCH_TRACE/tlparse 便于检查生成的代码,附上错生成的代码通常足以诊断,无需完整复现。
绕开编译器崩溃
对某些架构,模型的特定部分极难编译。可显式禁用问题部分、对能work的部分启用 PT2。三种做法:
- 把 torch.compile 移进模型内部,只编译特定子模块而非顶层整体。例如 PT2 对 DDP/FSDP 这类分布式 wrapper 支持不佳,就对传入 wrapper 的内部模块做 torch.compile;再如顶层全量编译不现实时,可只编译 transformer block 这类局部模块并照样获益。
- 继续在顶层 torch.compile,用
@torch._dynamo.disable()装饰器禁用模型特定部分。语义上这等价于只编译部分模块(disable 无非是在到达被禁模块时制造一个 graph break),但"声明什么不可编译"往往比"声明什么可编译"更方便。示例:推荐模型中稀疏架构难以编译,就用它在稀疏部分禁用 PT2。 - 某些配置更易诱发崩溃:用
dynamic=True崩溃时先移除该 flag(依赖 automatic dynamic 自行判断是否需要动态)。
想"头铁硬闯"(YOLO),可设 torch._dynamo.config.suppress_errors = True——编译器每次崩溃就跳过该帧稍后再试。该配置定义于 torch/_dynamo/config.py,默认取环境变量 TORCHDYNAMO_SUPPRESS_ERRORS。不过最好还是最终手动补上必要的 disable 注解。
在编译期打印信息(comptime)
正常代码里的 print 会引发 graph break。但若你关心的是编译期可知的量,可用 comptime.print 在编译发生时打印——比如查看某个符号对应的符号表达式,或 Dynamo 认为某个变量的类型/尺寸是什么:
from torch._dynamo.comptime import comptime
# ... 在你的代码中 ...
comptime.print(...anything...)
实现位于 torch/_dynamo/comptime.py:ComptimeVar.print(L237)在编译时求值并输出,另有 print_graph、print_disas、print_locals、print_guards、print_value_stack 等姊妹 API(L218-L398),是了解编译器视角的利器。
首要判断:是重编译太多,还是编译器本身就慢?前者可用 TORCH_TRACE/tlparse 或 TORCH_LOGS=recompiles 发现(通常伴随 cache_size_limit 超限的 warning)。
重编译过多(Compiler is recompiling too much)
tlparse 中某帧如出现 [10/0] [10/1] [10/2] ... 一直递增到 7(或你配置的最大 cache 数),基本可断定出了问题。默认到达帧缓存上限后会停止编译该帧,但仍会递归尝试其内部帧——解决编译时间问题最好先修最外层帧。第一道防线是用 TORCH_LOGS=recompiles 查看重编译原因(这些日志会告诉你 Dynamo 首次编译时 specialized 了什么、后续运行不再成立从而触发重编译)。一个注意点:当前 recompiles 日志只报告第一个失败的 guard,修掉后可能冒出另一个同样导致 specialization 的问题;可 diff 两次编译生成的 guards(用 TORCH_LOGS=guards 或 tlparse)来感知全貌。
尺寸相关的重编译
由尺寸变化触发时 guard 形如 L['x'].size()[0] == 2954。最常见原因是某处强制了 specialization。若被 specialized 成具体整数,可在该帧 compilation_metrics 的 "Symbolic shape specializations" 一节找到原因;否则用 TORCH_LOGS=dynamic 定位 guard(知道精确 guard 串时也可用 TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="s0 == 2" 之类只输出该 guard 的详细信息)。最常见的三种原因:
- 首帧按静态尺寸编译。默认首次编译假定所有尺寸静态;只有发现尺寸变化才重编译并尝试动态化。现象是经典的两次编译
[0/0](static)与[0/1](dynamic);若很多输入尺寸需要动态、但其中一些直到更晚的输入才变化,可能产生更多次重编译。对策:用torch._dynamo.mark_dynamic(tensor, dim)标记某维动态,立即以动态方式编译。但输入尺寸为 0/1 时该方法不适用,须用torch._dynamo.mark_unbacked(tensor, dim)——这会引入另一种报错:此后一旦对该尺寸量做 guard 就会抛 GuardOnDataDependentSymNode 错误(可能正是你想要的,也可能不是)。mark_dynamic/mark_unbacked的实现与签名见 torch/_dynamo/decorators.py。两者在有 graph break 时都表现不佳:mark_dynamic 注解不会穿过 eager 代码传播,你必须准确标注每个 graph break 区域的所有输入——除非你是专家否则基本做不到。 - 框架代码无法以符号方式运行,被迫 specialization。例如算子 schema 未 SymInt 化、某算子的 C++ meta 实现只接受
int64_t而非 SymInt,或代码写得糟糕导致 specialization。若能在 issue 中附上对应 specialization 的日志行(默认会给出触发 specialization 的框架代码最佳猜测),通常很快能修。更隐蔽的情形是重编译源于非常规 guard(如s0 % 2 == 0),此时需用TORCH_LOGS=dynamic追根溯源。 - 用户代码本身存在根本性障碍。如用整数索引 Python list 会强制 specialization(PT2 不支持对 Python 数据结构的动态访问)。这类情况可能需要相当深度的架构调整。
NN module(ID_MATCH)相关重编译
非常常见的重编译来源:对某对象装了 ID_MATCH guard,而该对象每轮迭代不断被重建——通常是 NN module(默认对 NN module 装 ID_MATCH guard)。对策:设置 torch._dynamo.config.inline_inbuilt_nn_modules = True 强制 Dynamo trace 进 NN module 内部,用更细粒度的 guard 取代 ID_MATCH guard。注意演进:原文写作时这仍是计划默认开启的选项;而当前仓库中它已是默认 True 且标记为 deprecated("does not do anything, inline_inbuilt_nn_modules is always True"),见 torch/_dynamo/config.py。
通用排查建议:读懂 guards
若重编译不属于上述常见原因,就要弄清 guard 从哪来。TORCH_LOGS=guards 输出(tlparse 中的 dynamo_cpp_guards_str 产物)会为每个 guard 标注来源,例如:
ID_MATCH: ___check_obj_id(G['g'], 7665376) # if g: # b.py:7 in f
此处 G['g'] 表示名为 g 的全局(局部变量则显示为 L['varname']),ID_MATCH 是 Dynamo trace 到 b.py 第 7 行时安装的——该行是对 g 的 if 判断,于是 guard 只验证 g 为真,完全合理(需要知道 g 的值才能决定分支走向)。
guards 输出按树组织(某些 guard 需先于其他 guard 求值才安全),但通常自解释。更完整的 guard 树示例(叶子即真正执行的逻辑):
TREE_GUARD_MANAGER:
+- RootGuardManager
| +- DEFAULT_DEVICE: utils_device.CURRENT_DEVICE == None
| +- GLOBAL_STATE: ___check_global_state()
| +- GuardManager: source=L['x'], accessed_by=DictGetItemGuardAccessor(x)
| | +- TENSOR_MATCH: check_tensor(L['x'], Tensor, DispatchKeySet(...), torch.float32, ...)
| | +- NO_HASATTR: hasattr(L['x'], '_dynamo_dynamic_indices') == False
| +- GuardManager: source=G, accessed_by=GlobalsGuardAccessor
| | +- GuardManager: source=G['g'], accessed_by=DictGetItemGuardAccessor(g)
| | | +- ID_MATCH: ___check_obj_id(G['g'], 7665376)
+- LAMBDA_GUARD: 2 <= L['x'].size()[0]
另一招:取同帧两次编译的 dynamo_cpp_guards_str 做文本 diff。guards 差异处就是两次编译为何不同的解释——与 TORCH_LOGS=recompiles 不同,这给出全部差异 guards 而非仅第一个。
找到触发 guard 的代码后,尝试改写为不触发 guard 的写法;若你认为 PyTorch 本应支持你的代码却加了不必要的 guard(over-guarding),则提交 bug。牢记本质:Dynamo 只能处理直线(straight-line)trace,若程序有分支且两个分支都要编译,就必须编译两次——除非你用 torch.cond 或其他高阶算子改写程序。
编译器整体偏慢
循环会被展开(Loops are unrolled!)
与传统编译器不同,torch.compile 的耗时取决于代码执行 trace 的长度,而非输入 Python 程序的规模。例如循环遍历 size 1024 的维度,就会得到 1024 份循环体副本——编译很慢,而且 PT2 不做自动向量化,结果也不会快。对策视情况而定:
- PyTorch 配合 einops 与 broadcasting 表达力很强,尽量改写成 batch 形式;难以从单 batch 改写为多 batch 时,
functorch.vmap能帮你以单 batch 编写后自动 batch 化。 - 也可用 Numba 这类"把循环代码变成优化 kernel"的库。
注意编译器的常数因子远差于 eager:eager 下性能平庸但足够原型验证的代码,torch.compile 后可能慢到不可接受。文档举例:约 1500 个输入 Tensor(每个代表一个稀疏特征)的程序,图里充满海量微型 tensor 计算时很容易陷入数小时的编译。
用缓存加速编译
缓存是降低编译开销(尤其反复运行同一模型)的重要投入方向。Meta 截至 2024-07 已部署远程 Inductor 缓存,训练/推理 warm start 时间分别改进约 50%/80%。缓存运维上有一定复杂度(参见"NCCL 超时"中 stuck ranks 一节),但值得部署。H2 2024 的优先级之一是把缓存扩展到 AOTAutograd,进一步削减编译耗时。
默认还存在一份文件系统缓存,保存在 /tmp。若系统对 /tmp 的保留时间不够长,用 TORCHINDUCTOR_CACHE_DIR 更改缓存目录。仓库中该环境变量的解析位于 torch/_inductor/runtime/cache_dir_utils.py,并被 torch/_inductor/codecache.py、torch/_inductor/async_compile.py 等模块在子进程/编译环境中显式传递。关闭缓存的开关则见 torch/_inductor/config.py(force_disable_caches,环境变量 TORCHINDUCTOR_FORCE_DISABLE_CACHES=1)与 torch/_inductor/config.py(远程缓存开关 fx_graph_remote_cache,环境变量 TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE=0)。
用分层编译(hierarchical compilation)加速
默认 PT2 把所有模型代码内联进单个函数再编译。对某些架构,被复用的块(如多个 transformer block 副本)会被重复内联编译。若你并不需要跨块边界融合带来的收益,可以只对 block 本身做 torch.compile——只编译一次、处处复用,大幅削减编译时间。此时可能需要 inline_inbuilt_nn_modules = True 确保 self 实例变化时不重编译(当前仓库该选项已默认开启)。H2 2024 计划支持在更大的 torch.compile 区域内将块 noinline,实现真正的分层编译。
剖析编译器自身
让编译器更快很难,但有时存在巨大且明显的低效点。用 TORCH_COMPILE_CPROFILE=1 可方便地拿到编译过程的 cProfile trace(会生成 dot 图,热点越红越热),采样型 profiler(如 py-spy)也能帮助定位——例如可发现大量时间花在 fork 子进程上。
编译结果首次运行慢、之后正常
PT2 的大部分工作在编译期完成,但有部分编译步骤被推迟到编译产物首次真正运行时(当前 tlparse 对这些步骤记账不佳,可用 TORCH_LOGS=inductor 获取部分信息)。最典型的是 CUDA graphs(mode="reduce-overhead" 或 mode="max-autotune" 激活)需要正式录制前的 warmup,产生特征性的三段式耗时曲线:
- 首次运行(约 100s):完成 PT2 编译,但以无 CUDA graphs 方式直接运行结果;
- 第二次运行(约 10s):无需 PT2 编译,执行 CUDA graph 录制;
- 第三次运行(<1s):真正的快路径。
此外,单个编译产物可能录制多个 CUDA graph:动态 shape 模型会被编译一次但为每个尺寸分别录图;同一 nn.Module 被多份不同参数复用时也会按参数集重录。录制过程可能比常规推理慢一个数量级。删除 mode 参数或改用 mode="max-autotune-no-cudagraphs" 可消除该 warmup 开销(代价是运行时变慢)。
从这里开始:实用技巧清单
文档要求先读这一节:
- 网络可接受降低矩阵乘法精度时务必开启 TF32——naive 基准测试发现 PyTorch 比别家慢的第一大原因就是没开 TF32。
- Inductor 针对 A100/H100 优化;连 V100 都已不在"快乐路径"上,消费级显卡尤其容易出问题(大量归因于 Triton 目标硬件集合有限)。社区欢迎外部贡献改善消费级显卡。
- 基准测试前务必 CUDA 同步再计时,否则会低估 kernel 耗时;同时保证至少两轮 warmup(用 CUDA graphs 时最好十轮左右)。
- 测静态 shape 性能时显式
torch.compile(dynamic=False),否则首个尺寸很快、后续尺寸被自动重编译为动态而变慢。 - 网络受 CPU 开销约束时考虑
mode="reduce-overhead"(启用 CUDA graphs)。 - 能等待编译时间时,
mode="max-autotune"会花时间自动调优模板与 block size,通常带来若干百分点收益;调优结果可缓存,不必每次重复。 - PT2 与手写自定义 kernel 配合良好;若算子对输入 stride 有严格要求,可考虑设置
needs_fixed_stride_order。
仓库侧佐证:autotune 相关配置(max_autotune、max_autotune_pointwise、max_autotune_gemm 等开关,对应环境变量 TORCHINDUCTOR_MAX_AUTOTUNE*)均集中在 torch/_inductor/config.py。
graph break 太多
tlparse 输出中可看到 PT2 编译了多少张图,也可用 TORCH_LOGS=graph_breaks 或 profiler 识别。graph break 损害性能的原因:
- 图更小 → 融合机会更少、PT2 固定开销(guards、AOTAutograd 运行时代码)占比更高。
mode="reduce-overhead"的 CUDA graphs 尤其受损:Tensor 必须拷入/拷出固定的 CUDA graph 地址。 - 小图更易呈现动态行为(遇到不同尺寸),而动态 shape 的代码生成性能整体弱于静态 shape。
多数 graph break 只是 Dynamo 缺功能——请提交 bug(即便你能用简化 Python 代码绕过,也请照提,社区希望支持这些模式)。有几类高发场景需要特殊处理:
- 任何 data-dependent 计算(布尔 mask、
item()、nonzero()等)默认触发 graph break。胆大者可设torch._dynamo.config.capture_scalar_outputs = True(对应环境变量TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1)与torch._dynamo.config.capture_dynamic_output_shape_ops = True,见 torch/_dynamo/config.py,但随后要准备应对 GuardOnDataDependentSymNode 错误。 - torch.compile 编译某帧失败后会再尝试其每个内层帧;若 graph break 深嵌在某个内层帧里,只会产生一堆碎小帧。不如在调用栈较高处
torch._dynamo.disable一个函数,避免 PT2 反复碾压大量小图。
如何把 kernel 映射回 Inductor 代码与图
Inductor 会尽量为 kernel 取与融合算子对应的有意义名称;torch/_inductor/config.py 中的 triton.descriptive_names 配置可调整命名规则。多数情况下可通过 tlparse 的 output_code 把 kernel 映射回原始用户代码。仍困难时可设 TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1(解析见 torch/_inductor/config.py)强制每个 kernel 唯一命名。
拿到生成的 Triton 代码后,每个 kernel 上方有回指源码的元信息:
# kernel path: /tmp/torchinductor_ezyang/tc/ctcg6vrb3wgwmuh625mqucd4gfgje4wzwya3gslwtfgvk74ucyl5.py
# Source Nodes: [gelu], Original ATen: [aten.gelu, aten.gelu_backward]
# gelu => add, erf, mul_1
Source Nodes:对应 dynamo_output_graph 中产生该 kernel 的节点名(tlparse 中可查);backward 生成的 kernel 可能为空。Original ATen:贡献该 kernel 的、decomposition 之前的原始 ATen 算子名。- 右侧
gelu => add, erf, mul_1:把 Dynamo 源节点映射到 inductor_post_grad_graph 节点。
后文 post_grad_graph 还会报告生成这些节点的用户代码源码。注意两点:
- kernel 与 Source Nodes 并非严格一一对应:单个源节点可能 lower 成多个 Inductor IR 节点并分别调度成 kernel(该节点会在多个 kernel 中出现);重计算(尤其在 backward 中重算 forward 算子)也会导致节点多次出现。
- 截至 2024-07,PT2 不打印 backward 对应的 post grad graph 节点(issue #130147),forward 的打印是准确的。
解读编译代码的 profile
对 torch.compile 代码可使用 PyTorch profiler——当模型很大且目标不明时很有用(能暴露 device-to-host 同步阻塞、提示该查哪些 kernel)。注意 torch.autograd.profiler.record_function 目前在 torch.compile 区域内不生效,可在其外部使用。
一个算子被启动时,profile 中通常可见三类事件:CPU 侧事件 → kernel launch(GPU kernel)→ GPU 侧事件。文档对比了各来源 kernel 的形态差异:
- Inductor 生成的 Triton kernel:(1) CPU 侧事件以
triton_前缀出现,当前信息较少(仅 kernel 名与 launch,不及典型 aten kernel 带 shape/type);(2) kernel launch 表现为cuLaunchKernel而非 aten 常见的cudaLaunchKernel;(3) GPU 侧事件名称的描述性取决于unique_kernel_names配置。 - 非 Inductor 生成的手写 Triton kernel:CPU 侧事件可能完全不出现(自动插入 profiler 事件的机制实现在 Inductor 层,绕开 Inductor 的 Triton kernel 除非手动标注否则不出现);launch 同样为
cuLaunchKernel;GPU 侧事件存在且命名接近你写的 Triton kernel。 - Inductor 生成的 CPU kernel:CPU 侧事件不出现(尚未加 profiling),也自然没有 launch 与 GPU 侧事件。
关于 Torch-Compiled Region:这是 Dynamo 为每个经手的帧添加的 profiler 事件。为什么图里有很多嵌套的 Torch-Compiled Region?因为发生 graph break 时,Dynamo 处理的帧会递归调用 continuation 帧,从而形成嵌套。两个各自独立 torch.compile() 的函数一般产生**相邻(不嵌套)**的 Torch-Compiled Region;而遇到 graph break(或 disable()/skip 区域)时则期望看到嵌套事件。
autograd 与 torch.compile:当图中有输入 requires_grad=True 时,profile 中会插入名为 CompiledFunction 的事件。它是 PT2 编译器实现中用于把编译后的 forward 与 backward 拼接起来的 autograd.Function(相关运行时 wrapper 逻辑在 torch/_functorch/_aot_autograd/runtime_wrappers.py 中;仅当部分输入需要梯度时才使用,否则编译器采用不含该事件的 wrapper)。backward 中通常配对一个 CompiledFunctionBackward 事件,两者间会出现 fwd-bwd link。不要求梯度的图不能用 CompiledFunction 判断是否编译成功——应看 Torch-Compiled Region 或 Inductor Triton kernel 是否存在。
DDP Optimizer:若启用,它会引入不显示为嵌套 Torch-Compiled Region 的 graph break——因为 DDP Optimizer 引入的 break 发生在 Dynamo 处理之后。
能否从 profile 数 graph break?忽略有多段独立 torch.compile 代码的情况,"Torch-Compiled Region 数量减 1"(即嵌套区域数)大致准确,但更推荐用 tlparse。
为何 Torch-Compiled Region 下有时看到 ATen 算子?多数 aten op 会被 Inductor lower 成 Triton kernel;但不常见的算子(及多数自定义算子)没有 Inductor lowering,会回退到原始 aten 实现。matmul 频繁回退 aten 有两个原因:gemm 未开 autotuning 时 Inductor 自动回退 ATen;autotuning 时也可能发现原始 ATen 实现仍比 Triton kernel 快而选用之。这不是 graph break,只是 Inductor 选择最优代码。
处理融合问题
可视化融合决策
看生成 kernel 的注释可以反推哪些 ATen 节点被融合,但 inductor_post_grad_graph 的图形化表示能更清楚地展现 kernel 间数据依赖:设 TORCH_COMPILE_DEBUG=1 INDUCTOR_ORIG_FX_SVG=1 会生成 SVG,每个 ATen 节点入图并按融合的 kernel 分组。真实模型上该图可能大到难以处理,最适用于你要理解编译行为的小代码片段。
为什么 Inductor 那样融合?
诊断性能时常见疑问是"为什么预期合成单个 kernel 的代码被编成了多个"。关键诊断手段是 TORCH_LOGS=fusion,它打印两个 buffer 为什么未被融合。要读懂它先要理解 Inductor 可期望执行哪几类融合(文档整理):
- Pointwise 融合:同一 shape 上多个 pointwise 操作会被融合;该融合在 lowering 而非 scheduling 阶段即时发生(表现为同一 Pointwise IR 节点)。
- View 消除:view 一般无需生成代码,只需改写后续访问的索引表达式,直接索引原始输出。
- 模式匹配(Pattern matching):一批 FX pass 寻找若干 FX 节点组合并用单一算子替换。
- Epilogue 融合:对拥有 Triton template 的 kernel(如 matmul),可把额外 pointwise 操作融合进其 epilogue。并非总是划算——复杂的矩阵乘常被精细调优,加算会使其劣于未融合版本。
- 垂直融合:若消费者的所有读取要么匹配生产者的对应写入、要么由可在两者融合前调度的节点写入,则可把消费者融入生产者。这类决策记录在
TORCH_LOGS=fusion。 - 水平融合:互不依赖但共享读取的两个节点也可融合,决策同样记录在
TORCH_LOGS=fusion。
不融合的原因很多;即便可融合,也可能存在多种互斥融合需 Inductor 取舍。以 TORCH_LOGS=fusion 为准。Triton 代码生成的主算法 can_fuse 位于 torch/_inductor/codegen/simd.py(放这里以便其他 Triton-like 后端复用)。要点:水平融合要求节点有兼容的 tiling。文档给出三个真实案例:
- float8 Linear 中
max(abs(tensor))与to_fp8_cast(与 max-abs 无数据依赖)本应融合却未融合,log 显示"invalid tiling for reduction":pointwise 的 tiling 为 (4096, 4096),reduction 的 tiling 为 (512, 32768)。根因是 [4096×4096]→[1] 的 amax 为效率被拆成两层归约;Inductor 选择第一层中间大小 X 时没意识到取 4096 才能与后续 pointwise 融合。有趣的是,最终结论是这两者融合其实并非好事(第二层直接从 4096 归约到 1 反而低效);有 issue 跟踪让 Inductor 把循环顺序决策推迟到融合之后。 - 有用户在 profile 中发现 MHA 的 QKV projection 后 add/view Triton kernel 耗时且明显受内存带宽约束,期望 add 被融进 mm 变成 addmm;但该自定义 MHA 中 mm 与 add 之间隔着 view,阻碍了融合。把 mm → view → add 重排为 addmm → view 的 POC 显著改善性能。长期方案是研究让 view 尽量移出计算算子路径的优化 pass。
- 有用户期望 pointwise 与 amax 融合,微基准中发现归约范围含多个动态维时索引表达式过于复杂而无法融合;仅单个动态维(如 batch 维)时可融合。
原则上更多 autotuning 能帮 Inductor 减少启发式依赖、尝试各种融合。但当前 autotuning 实际上只用于 Triton block size 决策。
调优 Inductor 生成的 kernel
Inductor 尽力用 autotuning 自动选好 block size,但偶尔仍做出坏决策。专家用户常把 Inductor 当作进一步手调 Triton 实现的基础。可用工具:
- 生成的 kernel 自带 benchmark harness(
benchmark_kernel),直接运行 inductor_output_code Python 文件即可调用。tlparse 输出的 Python 文件可直接运行,作为优化的起点。 TORCHINDUCTOR_PROFILE=1让 Inductor 估算 kernel 带宽数字(配置见 torch/_inductor/config.py,另有输出文件与环境变量可调)。- ncu(Nsight Compute)是判断 kernel 低效与否的利器:例如
ncu_mem_bw_gbps低说明未饱和内存带宽;kernel 本身耗时长则是计算优化的候选。注意 ncu 数字比 Inductor 估算更真实(代表真实内存带宽)。文档案例:用该方法定位到tl.rand慢——它生成四个样本却丢掉三个。
直接读 kernel 也常有收获:文档描述某 bug 的排查正始于用户查看生成代码、发现 Inductor 生成非合并写入,对比了两种 tl.store 的索引模式(一种为 out_ptr2 + (r1 + (1024*x0)) 的非合并写法,另一种为换用 (x0 + (26624*r1)) 的合并写法)。
应对 GPU 空档(dead time)
GPU 利用率常被当作性能代理指标:GPU 计算出现大段空隙就是坏信号。未经 torch.compile 已利用良好的代码,编译后反而出现空隙,是因为 Inductor 优化过的 kernel 更快、花在内存带宽的时间更短,关键路径随之改变。
GPU 空档通常源于 DtoH 同步或分布式集合通信;也可能是大量 CPU 开销(但 PT2 总体削减 CPU 开销,mode="reduce-overhead" 可进一步降低)。消除 DtoH 同步 bubble 的最佳手段是消灭 DtoH 同步——例如用 boolean mask 时改考虑 torch.where。有些空隙无法避免(必须做集合通信、必须检查数据以削减下游计算量),此时通常的思路是找到不依赖阻塞操作的其他计算并提前调度,让它在等待通信时执行。原则上 PT2 能做这类优化,截至 2024-07 尚未实现,但社区对该类优化兴趣浓厚。
自定义 FX pass
单个模型中的某段低效代码或许能靠改写解决;但大量模型有同一问题时,可写一个自定义图 pass来识别并替换这些模式。注册点位于 torch/_inductor/config.py:
# 注册自定义图优化 pass 钩子。pre/post pass 分别只在 post_grad_passes
# 中 pattern_matcher 之前/之后应用。
# def my_custom_pre_pass(graph: torch.fx.Graph):
# # 我的自定义图优化 pass
# ...
# def my_custom_post_pass(graph: torch.fx.Graph):
# # 我的自定义图优化 pass
# ...
torch._inductor.config.post_grad_custom_pre_pass = my_custom_pre_pass
torch._inductor.config.post_grad_custom_post_pass = my_custom_post_pass
另有 joint graph(前向+反向联合图)的注册点 joint_custom_pre_pass / joint_custom_post_pass。这些钩子被实际调用的位置分别在 torch/_inductor/fx_passes/post_grad.py(post-grad 图)与 torch/_inductor/fx_passes/joint_graph.py(joint 图),缓存键中也包含这些自定义 pass(见 torch/_inductor/codecache.py),保证自定义 pass 变化会正确失效缓存。如何编写 FX pass 可参考仓库中大量既有 pass 实现(torch/_inductor/fx_passes/ 目录),你看到的算子取决于 pass 执行阶段,decomposition 之后一般是 ATen 操作。
CUDA graphs 专项建议(mode="reduce-overhead")
必须理解 CUDA graphs 的基本特性:
- 只能捕获 CUDA 计算。图内有 CPU 计算就无法用 CUDA graphs;有时 CPU 计算可搬到 GPU(Inductor 原则上可自动做,但只覆盖有限场景)。随机数可与 CUDA graphs 协作,但 checkpointing 与 RNG 不行(issue #130123 跟踪,另有 PR 提出修复策略)。
- 录制好的 CUDA graph 运行在固定 CUDA 地址上。最坏情况下输入/输出必须拷入/拷出这些固定地址。实际有一系列缓解优化:参数按静态地址处理、假定不移动;输出默认留在 CUDA graph 内存(重跑前必须先释放)。违反这些约束会导致重录图(如参数地址变化)或额外拷贝——内存确实从不移动时,可用
torch._dynamo.mark_static_address避免拷贝。 - 录制的 CUDA graph 必须保活其分配的内存(只作用于特定地址),因此即使当前未使用,活跃 CUDA graph 的内存也无法被 eager 复用,总体可能抬高常驻内存。graph break 间的多个 CUDA graph 可通过 CUDA graph trees 复用内存。
- CUDA graphs 不支持动态 shape。
dynamic=True与mode="reduce-overhead"虽可并存,但实现方式是为每个见过的尺寸各录制一张 CUDA graph。尺寸太多时可用 padding 到倍数来减少需录制的图数量(另见"首次运行慢"小节)。 - 若报错
"Exception Found: These live storage data ptrs are in the cudagraph pool but not accounted for as an output of cudagraph trees",可尝试torch._inductor.config.triton.cudagraph_trees_history_recording = True(配置定义于 torch/_inductor/config.py,在 torch/_inductor/cudagraph_trees.py 中用于输出内存分配来源提示)。
最常见的抱怨是 CUDA graphs "没生效"。用 TORCH_LOGS=perf_hints 可得到未生效原因的解释。常见补救是把模型拆成"可 CUDA graph 化"与"不可 CUDA graph 化"两部分。
torch.compile 通常预期降低内存(除非用 mode="reduce-overhead",CUDA graphs 会增加内存:图内存开销 + 输入拷入图内存)。但 PT2 的 bug 也可能导致内存超标,历史案例:
- Inductor 按对齐输入编译图时,若传入未对齐输入就必须拷贝对齐,该拷贝抬高总内存。
- 不再使用的 Tensor 仍被保活。Python 中一个顽固来源:Tensor 作为函数参数时,默认在整个函数调用期间保活,即使中途已不再使用。PT2 有 boxed calling convention(传入参数列表并在函数内清空)绕开该问题,但偶尔实现有误。
- 坏融合决策重排节点导致高水位内存上升。Inductor 本应不做会抬高内存高水位的融合。
内存剖析
可用 memory snapshots 可视化 PyTorch 的已分配内存状态(PT2 下可用)。这对分析 CUDA graphs 内存尤其有用。
##(选择性)Activation checkpointing
Activation checkpointing 是经典的"以算换存"。标准 checkpointing API 与 torch.compile 兼容;文档给出一份训练吞吐最大化案例研究(利用 checkpointing 压缩激活内存、扩大 batch 提升吞吐)。
一般而言,NCCL 超时源于并非所有节点都及时执行同一个 NCCL 集合。即使没有 PT2,NCCL 超时也是分布式训练的梦魇——诱因可能是间歇性硬件问题、网络交换机故障、部分节点非确定性地触发 Python GC、统计 profiler 只采样到 rank 0、框架 bug……而 PT2 在这一切之上又引入了新奇的失败模式。应对 NCCL 超时最重要的诊断是检查挂起时各 rank 的栈——重点是那些没赶到 NCCL 集合的 rank 在做什么。下文按卡住 rank 的行为分类。
卡住的 rank 在编译
PT2 存在已知问题:模型编译可能远超 NCCL 超时阈值。若只有部分 rank 在编译而其他没有,或部分 rank 缓存命中而其他没有,编译耗时差极易触发 NCCL 超时(仓库 issue #108971 计划修复)。等待修复期间的缓解手段:
- 调大 NCCL 超时以覆盖编译时间。不算好方案(真出问题时你也要等这么久才超时),但假设没有其他问题,可帮你解除阻塞、确认"编译耗时就是元凶"。若分歧只发生在前 N 轮,可 warmup 期临时用长超时、稳定后再调回合理值。
- 使用远程缓存时的 rank 失衡:部分 rank 命中远程缓存、部分未命中就会失衡(Meta 内部在 pyper 模型上观察到)。此时禁用远程缓存可让所有 rank 同步编译、同步抵达集合通信。仅关远程缓存:
TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE=0或torch._inductor.config.fx_graph_remote_cache = False;本地全关缓存:TORCHINDUCTOR_FORCE_DISABLE_CACHES=1或torch._inductor.config.force_disable_caches = True。 - 某 rank 在重编译时的合理策略是从源头阻止重编译——让每个 rank 只在训练开始编译一次、之后绝不重编译。参考"重编译过多"一节。
卡住的 rank 在运行编译代码
PT2 可能引入的理论性问题:
- PT2 会改变代码的性能特性。若性能变化在 rank 间出现分歧(如某种 data-dependent autotuning),数据依赖的性能差异可能放大到触发 NCCL 超时——此时调大超时是合理的。排查方法与无 torch.compile 时诊断不平衡一致:查看 GPU trace,看是否存在不平衡。注意:当编译产物是动态 shape 时,不同 rank 各自为不同尺寸录 CUDA graph 也可能放大分歧(参见 CUDA graphs 一节)。
- PT2 支持捕获分布式集合通信,也就有机会优化它们(如重排)。若 PT2 重排了集合,必须让所有 rank 以完全相同方式重排;一旦该过程不确定、或各 rank 输入图不同,就没有办法验证 PT2 做得正确(当前编译期不做跨节点通信)。
- 配置项
torch._functorch.config.fake_tensor_propagate_real_tensors可能造成挂起——传播真实 Tensor 需要在编译期运行真实算子,包括集合通信(issue #126846)。
文档坦言:这一块的调试手段目前还比较欠缺——先做消融(见前文),祝好运。一个较大的空白是"在模块间记录中间值、以此三角定位与 eager 的分歧点"的工具,OSS 中尚不存在。
若干让 Dynamo 更"较真"的配置值得一试(仓库源码均已核实):
:大量与 NN module 的 eval/train 切换、或两次编译调用间 module 突变相关的精度错误,源于对 NN module 缺 guard。此配置现已默认开启(OSS 自 2024-07-03、fbcode 自 2024-09-27),已无需手动设置。torch._dynamo.config.guard_nn_modules = Truetorch._inductor.config.emulate_precision_casts = True:编译低精度代码时,若某次 downcast/upcast 非必要(例如两个算子已融合),我们会省略到小类型的转换——这会改变数值行为,某些场景下高精度计算反而损害整体精度。该开关强制精确模拟精度转换(即使拖慢 kernel)。源码注释见 torch/_inductor/config.py:eager 计算 bf16/fp16 是升到 fp32 计算后再降回;两个低精度算子融合时 Inductor 省略中间的降/升对,通常其行为更接近 fp64 参考数值,但此开关可保留降/升对以完全模拟 eager 数值。torch._dynamo.config.optimize_ddp = False:DDP optimizer 相当复杂,值得先排除——尤其当你在backend="eager"下失败时。演进说明:当前仓库中该配置不仅接受布尔值,还扩展为True / False / "ddp_optimizer" / "python_reducer" / "python_reducer_without_compiled_forward" / "no_optimization"的细粒度枚举(默认True,语义为启用 ddp_optimizer),见 torch/_dynamo/config.py。
Stride 分歧(stride divergence)
PT2 不保证产出与 eager 完全同 stride 的 Tensor。若 torch.compile 区域之外的代码期望特定 stride 却不做检查、或使用了对输入 stride 敏感的受限 API(主要是 torch.as_strided,以及通过 .reshape()/.contiguous() 就地修改的冷僻情形),就可能引发正确性问题或运行时崩溃。
文档维护方持续在版本库更新本手册,核心变更包括:
- 2024-09-27:精度小节新增
torch._dynamo.config.optimize_ddp = False建议;guard_nn_modules因已默认开启而标记为过时。 - 2024-08-13:新增
torch._inductor.config.emulate_precision_casts = True与torch._inductor.config.triton.cudagraph_trees_history_recording = True两条说明。 - 2024-07-11:微基准 CUDA 代码时务必 CUDA 同步。
- 2024-07-10:CUDA graphs 需要两轮 warmup(而非一轮)。
- 2024-07-09:新增"Stride divergence"一节;澄清 compiled autograd 通过把 accumulate grad 节点融合进编译区域带来收益。
- 2024-07-08:本地缓存默认存
/tmp,可用TORCHINDUCTOR_CACHE_DIR更改。 - 2024-07-07:新增"What you should expect to compile"一节。
对照当前仓库可知,其中数条演进已被实现落地:inline_inbuilt_nn_modules 已默认开启并废弃(见 torch/_dynamo/config.py),optimize_ddp 从布尔扩展为枚举(torch/_dynamo/config.py),guard_nn_modules 相关行为已并入默认配置。这恰好印证了本手册的定位——它是一份随 torch.compile 快速迭代的"活文档",使用时应以你所安装 PyTorch 版本的实际行为为准。
延伸阅读建议(仓库内部):想深入了解编译器实现,可阅读 torch/_dynamo/config.py(Dynamo 全部配置项)、torch/_inductor/config.py(Inductor 全部配置项与 autotune/缓存/命名控制)、torch/_inductor/fx_passes/joint_graph.py 与 torch/_inductor/fx_passes/post_grad.py(自定义 pass 钩子调用点)、torch/_inductor/codegen/simd.py(Triton 融合决策算法 can_fuse)、torch/_dynamo/compiled_autograd.py(编译 autograd)、torch/_dynamo/comptime.py(编译期打印工具)以及 torch/_inductor/cudagraph_trees.py(CUDA graph 池管理)。
atomcodeClaude Code 的开源替代方案。连接任意大模型,编辑代码,运行命令,自动验证 — 全自动执行。用 Rust 构建,极致性能。 | An open-source alternative to Claude Code. Connect any LLM, edit code, run commands, and verify changes — autonomously. Built in Rust for speed. Get StartedRust0631
MiniCPM5-2BMiniCPM5-2B 是一款面向端侧、本地部署和资源受限场景的 2B 稠密 Transformer,能够达到同尺寸开源模型 SOTA 水平。Markdown00
video-shotcraftAI宣传片skill,使用 Remotion 制作电影级产品视频:提供106 张镜头配方卡和可复用的视频魔板。适用于 Claude Code 与 Codex以及所有其他智能体Markdown00
HivisionIDPhotos⚡️HivisionIDPhotos: a lightweight and efficient AI ID photos tools. 一个轻量级的AI证件照制作算法。Python09
DragonOSDragonOS is an operating system developed from scratch using Rust, with Linux compatibility. It is designed for **Serverless** scenarios. 使用Rust从0自研内核,具有Linux兼容性的操作系统,面向云计算Serverless场景而设计。Rust00
Spark-X2.5-1.7BSpark-X2.5-1.7B 旨在让强大的 AI 更加实用、高效且易于获取。这些模型在广泛的日常任务中表现出色,涵盖对话、写作、翻译、推理、编程、工具调用和智能体工作流,并在同等规模的开源模型中取得领先结果。Spark-X2.5 将面向效率的架构与最高 1M tokens 的原生上下文窗口相结合,并支持 200 多种语言。Python00