作者|吴育昕
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