如何看待PyTorch 2.0？

作者｜吴育昕
 

1

为什么是TorchDynamo
 

Graph capture 把用户 Python 写的模型代码变成 graph，是一切编译的根基。而 PyTorch 在试了这么多方案之后似乎已经锁定 TorchDynamo 作为 graph capture 的未来方向了，所以写一点关于 TorchDynamo 的内容，主要是解释到底为什么要做这个东西（离开FB一年了，内容主要凭自己的猜测和理解)。

一句话尽量解释 TorchDynamo 干了什么：利用 PEP523(https://peps.python.org/pep-0523/) 的 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 才能保证正确性。

而当 Torc