如何看待PyTorch2。0？

　　作者吴育昕
　　1hr为什么是TorchDynamo
　　Graphcapture把用户Python写的模型代码变成graph，是一切编译的根基。而PyTorch在试了这么多方案之后似乎已经锁定TorchDynamo作为graphcapture的未来方向了，所以写一点关于TorchDynamo的内容，主要是解释到底为什么要做这个东西（离开FB一年了，内容主要凭自己的猜测和理解）。
　　一句话尽量解释TorchDynamo干了什么：利用PEP523（https：peps。python。orgpep0523）的API在用户执行每个pythonframe前，拿到这个frame的bytecode，把其中认识的部分用tracing的方式提取出graph（并送给后端编译），不认识的部分维持原样。把修改后的bytecode还给CPython跑。
　　由于LazyTensor和TorchDynamo都做tracing，并且都是besteffortgraphcapture，即只编译自己能capture的部分，capture不到的用Python跑（akaPythonfallback），所以观感上两者可能会差不多。
　　然而，这两个方案的差别正是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代码都相当于不存在了，连symbolicexecution的overhead都没有，而被换成了编译后的nativecode。这一点在以前所有的partialgraphcapture的方案里是做不到的：
　　LazyTensor即使编译过的graph也要每次重新在Python里trace一遍，才能发现哦，这个graph我曾见过的。torch。jit。script、tf。function、jax。jit可以把装饰的pythoncode换成编译后的，但是这都依赖用户把这个subgraphrefactor出来放到一个单独的函数里。而TorchDynamo是全自动不需要用户改代码的。这种refactor除了增加额外的工作量之外，还可能与用户的代码结构冲突，因为用来编译的graph的边界与用户代码需要的抽象边界很可能不match：例如用户本来希望写三个函数，但是最佳的优化是把其中两个半函数变成一个graph，这会让用户很尴尬。
　　这只是一个最直接的例子。由于能够读写bytecode，理论上TorchDynamo能access更多LazyTensor根本没有的信息，做更多事情（后面会提到）。而读写bytecode的难度比sourcecode要低不少，所以成为了一个可行的方案。
　　2hrwholegraphcapture用处不大？
　　有的人可能会说，上面提到的东西对wholegraphcapture没太大用啊。
　　我觉得确实是这样：TorchDynamo是一个对partialgraphcapture追求极致的方案，能够对几乎所有的Python实现的模型开箱即用有加速，不用改代码前提是还要跑Python作为fallback。但是部署一般需要的是wholegraphcapture整个模型在一个graph里不能用Python。
　　用tracing做wholegraphcapture的前提是用户要在Python代码里避免所有不能被trace的东西，最常见的用户要做的三件事是：使用symbolicshape，使用symboliccontrolflow，禁用除了当前tensorlibrary之外的所有其它library。如果用户做到了这些，那只要一个普通的symbolictracing就能capture到完整的graph了，不需要TorchDynamo这么复杂的机制。TorchDynamo可能可以略微简化用户做这些的工作量，但我感觉不会有本质不同。
　　我个人的观点是，从实用角度出发，要求用户做上面几件事不算是太复杂的要求：禁用其他library理所应当就不说了；即使今天PyTorch还没有很好的symbolic｛shape，controlflow｝，但是只要用torch。jit。scriptiftracing来处理少量的symbolicshape和symboliccontrolflow，大多数模型都是可以正确的被torch。jit。tracecapture的。Meta应该有几十上百个vision模型实现在detectron2d2go里，目前基本都是走这条路部署的（我另有篇文章https：ppwwyyxx。comblog2022TorchScriptTracingvsScripting介绍这里面的细节）。
　　TensorFlow的wholegraphcapture就简单了：TF从第一天就有很好的symbolicshape和symboliccontrolflow，用就完了。tf。autograph甚至还自动化了一部分controlflow的改写工作。
　　所以，用户少量改代码仍然是必须的。当然，TorchDynamo毕竟有着改变用户要跑的bytecode的超能力。所以如果愿意的话，理论上可以让用户的wholegraphcapture工作变得更简单。例如：
　　模型中间的一些像ifx。shape〔0〕100的分支，有的可以通过shapeinference等价转移到模型开头的。这样的话就可以capture到更大的没有分支的subgraph。这件事在TorchDynamo里现在叫做guard。理论上可以把pythoncontrolflow自动替换成symbolic的，类似tf。autograph做的事情，只不过输入是bytecode而不是sourcecode。
　　目前TorchDynamo的nopython模式就是wholegraphcapture了。不过似乎还不是工作重心（以下引用自https：docs。google。comdocumentd1tlgPcR2YmC3PcQuYDPUORFmEaBPQEmo8dsh4eUjnlyIeditheadingh。rmxeybu31e0）：
　　PT2willprovideinfrastructureforanopythonexportmodeforedgeandperformancesensitiveservingcases。ThePT2teamwon’tdrivethisendtoendstack，butwewillkeepafeedbackloopwiththeteamsinchargeofthisandensurethecomponentswebuildarereusableinthesesituations。
　　不过与此同时，PyTorch2。0最近在完善symbolicshape的支持；functorch里也加入了少量controlflowoperator。这算是利好wholegraphcapture的消息。
　　3hr总结
　　总的来说，由于TorchDynamo在bytecode层面做文章，能做到一些其他方案做不到的事情。它的优点主要为partialgraphcapture服务：让用户的Python模型代码在0修改的情况下就能capture并获得加速。这体现了PyTorch对于Pythonfirst哲学的执念。这种执着是否有必要，见仁见智。
　　TorchDynamo的主要优势来自对bytecode的读写。JITscriptingcompiler的失败表明在sourcecodelevel做不了太多事，TorchDynamo能在bytecodelevel做事情确实很巧妙。不过，要完整的复刻CPythonbytecodeinterpreter，它的工作量、维护难度（以及出bug的概率）都是不小的。
　　另外，TorchDynamo对wholegraphcapture没有很大的帮助。对于复杂的模型，用户该做的改写还是得做。不过我估计2。0至少能对用户该做什么有个清晰的说法。
　　当然，最后PT2到底能不能把compiler做好，还有很多其他因素：IR怎么设计，何时specializerecompile，各种backend不同的特性等等。比如TorchDynamo和LazyTensor使用的IR其实也不一样。但是本文只讨论graphcapture，其他问题就不提了。
　　（本文经授权后发布。原文：https：www。zhihu。comquestion570220953answer2798657470）
　　欢迎Star、试用OneFlow最新版本：https：github。comOneflowInconeflow