PyTorch 2.0: TorchDynamo

Key Takeaways

do this [Music] I was waiting for all this talk for him to utter the word pie torch 2.0 and finally he said it welcome to pytorch conference this is my first one um I support the pi torch compiler team so as a compiler person PT 2.0 to a large extent is really a compeller story and you might wonder why does it take so long for us to actually fully Embrace compiler technology into the core of Pi torch and the reason is because what makes pie torch be loved by researchers is flexibility dynamism expressiveness is exactly what makes it very hard to compile so it took us five years of developing various Solutions compiler solutions to get to where we are and I want to call the inflection point of our journey searching for the right compiler solution polish Dynamo which is what I'm going to talk about so it starts with asking this question what is one thing that we could change first in the compiler stack that will fundamentally change the efficiency of Pi torch it's not a new compiler optimization it's not even a new compiler backend actually the first the key first move lies in the top part of the compiler which is the graph capture so we all know that graph mode execution can be a lot more efficient than executing Ops one at a time especially if you are a machine learning accelerator designer or if you are a production engineer unfortunately um for eager first machine learning framework like Pi torch it's actually non-trivial to get graphs reliably so that's why we care about graph capture because it's simple if you don't have reliable graph capture you don't have a reliable compiler coverage so let me introduce torch Dynamo our first out of box graph capture here's one example of a toy pytorch program and here is this one-liner change Torchlight compile and this in this example we actually allow you to specify a custom backend and in this particular example this my compiler basically takes the graph that's captured by Dynamo and do something about it which basically printed out and return a python callable so if you execute this program you'll be surprised you see actually three graphs being printed out because in this example we actually put a Twist of a data dependent control flow so this graph break actually is the feature it's a fundamental design feature of torch Dynamo that makes it out of box in sound so let me double click on what makes torch Dynamo sound and outer box number one is partial graph capture so this is our ability to skip anything unwanted and then fall back to eager and with partial graph capture we save the programmers to have to shoehorsing their models into what the back-end compiler want instead we'll just seem to be translating out in and out of between eager and the graph mode and then there is the guarded graph generation basically we want to make sure whatever we're capturing at capture time is still valid when we're replaying this graph and finally if we did find a mismatch on the guard we have the ability to just in time recapture them so basically partial guarded graph capture with just in-time recapturing is what set Dynamo apart from all our previous graph capture Technologies and by the way there's this thing called aot autograph that captures the backward part of the graph and that makes it working for training so I want to show you some numbers which assume is already highlighted in his keynote that would back the claims we made so far you have seen the 7K plus model so this demonstrates that it indeed works out of box and right now we have more than 20 experimental backends for inference that's integrated with torch Dynamo and we have more than one backhand that's integrated with torchamel for training and the number is smaller on trading side now because of Dynamo but because they're just less number of training back-ends available and we're working actively with vendors to migrate into this new stack and finally you've seen this 30 number what does it say actually it's not only talking about the performance it also demonstrates that partial graph actually works so by capturing crafts and allowing graph breaks we not only did not diminish the capability of the back end optimizing compiler we actually amplify it by bringing more models or more regions of the models into the compiler world so Dynamo is great but if you are an experienced pytorch compiler user you would ask this question so which technology do I use as of today we actually recommend using torture animal for any scenarios that can be deployed with eager so that includes most of the training scenario and some of the inference scenarios and for all the other scenarios which we call the exports past scenario our existing solution still works but in the future actually not so far future we actually want to have a con organic consolidation uh of front-end technology around torch Dynamo um it's it has always been our vision to have a unifying front end and have a vibrant thriving back-end ecosystem and the speed of our consolidation actually depends on two things number one how fast do we deliver the pytorch 2.0 export story and number two is how fast the vendors can migrate into this new stack so we look forward to that new experience of a Consolidated front end yeah so that's Dynamo and I want to bring you back to the question we asked in the beginning of the talk what is that one thing right so as a compiler optimizing compiler person for two decades and never imagine myself talk about a computer friend end but the front end part of the compiler graph capture really matters because once we're able to bring most of the pytorch models simply into the graph world the era of compiler accelerated pytorch really arrives so if you are excited about this Vision there are more Nuance aspects of the compiler story that you are going to hear today and in the following months and if you're a ml practitioner I would encourage you to try out pt2 on your workload and if you're a vendor I encourage you to talk to us to explore migrating into this new stack and finally we have some talk series that would dive deeper into the technology in in the pytorch 2.0 stack so with that I'm gonna pass the torch to Jason Enzo [Applause] thank you