作者|吴育昕
1
为什么是 TorchDynamo
Graph capture 把用户 Python 写的模型代码变成 graph，是所有编译的根基。而 PyTorch 在试了这么多计划之后仿佛曾经锁定 TorchDynamo 作为 graph capture 的将来方向了，所以写一点对于 TorchDynamo 的内容，次要是解释到底为什么要做这个货色（来到 FB 一年了，内容次要凭本人的猜想和了解 )。
一句话尽量解释 TorchDynamo 干了什么： 利用 PEP523(https://peps.python.org/pep-0…) 的 API 在用户执行每个 python frame 前，拿到这个 frame 的 bytecode，把其中意识的局部用 tracing 的形式提取出 graph (并送给后端编译)，不意识的局部维持原样。把批改后的 bytecode 还给 CPython 跑。
因为 LazyTensor 和 TorchDynamo 都做 tracing，并且都是 best-effort graph capture，即只编译本人能 capture 的局部，capture 不到的用 Python 跑 (aka Python fallback)，所以观感上两者可能会差不多。
然而，这两个计划的差异正是 TorchDynamo 要害的中央:
LazyTensor 是个纯靠 tracing 的计划，不可避免的问题是 「只能看见 trace 到的局部，只有 trace 一下才晓得哪里不能 trace」。而每次执行模型的时候，不能 trace 的局部可能不太一样。为了保障正确性，LazyTensor 就不得不每次执行都要从新 trace。举个极其的例子，模型里写了一个 torch.add(tensor, random.random())，其中 random 是个 LazyTensor 看不见摸不着的 Python 函数，那只有从新 trace 能力保障正确性。
而当 TorchDynamo 批改 bytecode 的时候，事件就不太一样了：
- 在 bytecode 里可能看得见所有须要的信息，所以可能证实「这段模型代码没有用到奇怪的货色所以不须要从新 trace」。
- 光证实了「不须要 trace」不代表能够真的不 trace，因为用户的代码还是一行行给 Python 来跑的。然而 TorchDynamo 又来了：CPython 到底跑什么 bytecode 是能够被它换掉的!
因而它能够做到这么一件事：当用户 call 一个被 capture 过的模型时，模型里大部分 Python 代码都相当于不存在了，连 symbolic execution 的 overhead 都没有，而被换成了编译后的 native code。这一点在以前所有的 partial graph capture 的计划里是做不到的：
- LazyTensor 即便编译过的 graph 也要每次从新在 Python 里 trace 一遍，能力发现「哦，这个 graph 我曾见过的」。
- @torch.jit.script、@tf.function、@jax.jit 能够把装璜的 python code 换成编译后的，然而这都依赖用户把这个 subgraph refactor 进去放到一个独自的函数里。而 TorchDynamo 是全自动不须要用户改代码的。
- 这种 refactor 除了减少额定的工作量之外，还可能与用户的代码结构冲突，因为「用来编译的 graph 的边界」与「用户代码须要的形象边界」很可能不 match：例如用户原本心愿写三个函数，然而最佳的优化是把其中两个半函数变成一个 graph，这会让用户很难堪。
这只是一个最间接的例子。因为可能读写 bytecode，实践上 TorchDynamo 能 access 更多 LazyTensor 基本没有的信息，做更多事件（前面会提到）。而读写 bytecode 的难度比 source code 要低不少，所以成为了一个可行的计划。
2
whole-graph capture 用途不大？
有的人可能会说，下面提到的货色对 whole-graph capture 没太大用啊。
我感觉的确是这样：TorchDynamo 是一个对 partial-graph capture 谋求极致的计划，可能对简直所有的 Python 实现的模型开箱即用有减速，不必改代码——前提是还要跑 Python 作为 fallback。然而部署个别须要的是 whole-graph capture 整个模型在一个 graph 里不能用 Python。
用 tracing 做 whole-graph capture 的前提是用户要在 Python 代码里防止所有不能被 trace 的货色，最常见的用户要做的三件事是：应用 symbolic shape，应用 symbolic control flow，禁用除了以后 tensor library 之外的所有其它 library。 如果用户做到了这些，那只有一个一般的 symbolic tracing 就能 capture 到残缺的 graph 了， 不须要 TorchDynamo 这么简单的机制。TorchDynamo 可能能够稍微简化用户做这些的工作量，但我感觉不会有实质不同。
我集体的观点是，从实用角度登程，要求用户做下面几件事不算是太简单的要求：禁用其余 library 理所应当就不说了；即便明天 PyTorch 还没有很好的 symbolic {shape, control flow}，然而只有用 @torch.jit.script_if_tracing 来解决大量的 symbolic shape 和 symbolic control flow，大多数模型都是能够正确的被 torch.jit.tracecapture 的。Meta 应该有几十上百个 vision 模型实现在 detectron2/d2go 里，目前根本都是走这条路部署的（我另有篇文章 https://ppwwyyxx.com/blog/202… 介绍这外面的细节）。
TensorFlow 的 whole-graph capture 就简略了：TF 从第一天就有很好的 symbolic shape 和 symbolic control flow，用就完了。tf.autograph 甚至还自动化了一部分 control flow 的改写工作。
所以，用户大量改代码依然是必须的。当然，TorchDynamo 毕竟有着 ” 扭转用户要跑的 bytecode” 的超能力。所以如果违心的话，实践上能够让用户的 whole-graph capture 工作变得更简略。例如:
模型两头的一些像 if x.shape[0] > 100 的分支，有的能够通过 shape inference 等价转移到模型结尾的。这样的话就能够 capture 到更大的没有分支的 subgraph。这件事在 TorchDynamo 里当初叫做 “guard”。
实践上能够把 python control flow 主动替换成 symbolic 的，相似 tf.autograph 做的事件，只不过输出是 bytecode 而不是 source code。
目前 TorchDynamo 的 “nopython” 模式就是 whole-graph capture 了。不过仿佛还不是工作重心 (以下援用自 https://docs.google.com/docum…）:
PT2 will provide infrastructure for a no python export mode for edge and performance sensitive serving cases. The PT2 team won’t drive this end to end stack, but we will keep a feedback loop with the teams in charge of this and ensure the components we build are reusable in these situations.
不过与此同时，PyTorch 2.0 最近在欠缺 symbolic shape 的反对；functorch 里也退出了大量 control flow operator。这算是利好 whole-graph capture 的音讯。
3
总结
总的来说，因为 TorchDynamo 在 bytecode 层面做文章，能做到一些其余计划做不到的事件。它的长处次要为 partial graph capture 服务: 让用户的 Python 模型代码在 0 批改的状况下就能 capture 并取得减速。这体现了 PyTorch 对于 “Python first” 哲学的执念。这种执着是否有必要，见仁见智。
TorchDynamo 的次要劣势来自对 bytecode 的读写。JIT scripting compiler 的失败表明在 source code level 做不了太多事，TorchDynamo 能在 bytecode level 做事件的确很奇妙。不过，要残缺的复刻 CPython bytecode interpreter，它的工作量、保护难度（以及出 bug 的概率）都是不小的。
另外，TorchDynamo 对 whole-graph capture 没有很大的帮忙。对于简单的模型，用户该做的改写还是得做。不过我预计 2.0 至多能对「用户该做什么」有个清晰的说法。
当然，最初 PT2 到底能不能把 compiler 做好，还有很多其余因素：IR 怎么设计，何时 specialize/recompile，各种 backend 不同的个性等等。比方 TorchDynamo 和 LazyTensor 应用的 IR 其实也不一样。然而本文只探讨 graph capture，其余问题就不提了。
（本文经受权后公布。原文：https://www.zhihu.com/questio…）
欢送 Star、试用 OneFlow 最新版本：https://github.com/Oneflow-In…