Model-Changing Transforms With Torch.compile - Thomas Viehmann, Lightning AI

In this talk, Tom Vman discusses model changing transforms using Torch Compile, highlighting the benefits of optimizing PyTorch models for better performance. He explains how optimizations can be composed for different models and configurations, and shares lessons learned from working with Lightning AI and NVIDIA's optimization tools, such as Lightning Thunder. Vman emphasizes the importance of creating flexible and composable transforms that allow developers to seamlessly modify and optimize their models. He also delves into practical aspects of using these transforms, including handling state dictionaries and optimizing model execution for various hardware configurations. The session showcases the potential of innovative techniques to streamline the development process for both novice and experienced PyTorch users.

Hello, I'm uh Tom Vman. I work with Lightning AI uh and have been doing PyTorch for a while. So, it's always a great pleasure to meet some old friends, meet some new friends. And uh today I also have the had the opportunity to bring two of my kids so they see their dad doing something different for work than sitting on the couch all day. So, uh awesome

to be here. Uh thanks for having me. Uh what I want to talk about is model changing transforms with torch compile which is kind of like exotic and so the first part will be why we're doing that. But first you need need to sit through some advertisement and with lightning AI we have the AI cloud that PyTorch developers love. Um, and you will see like I'm using

it uh and lots of others and you should try too and uh go check it out and get some free credits to give it a spin. And most of the people that try it love it. And you also get pre-made uh studios which are pre-made environments where you can clone it and edit the code, play with it uh all uh from from a great start. uh and

then you can when you're very happy you can also deploy your models or uh access the AI APIs of of the day through our platform and we're very proud to uh have very good and performant support for it. Okay, so part of that is making things fast and how do we do that? So the goal here in the talk is to remove friction in going from unoptimized

to optimized. And so this might not be all magic like we love the torch compile because it's all magic and you put something in and it gets passed out but sometimes you need to like do other things but it should be efficient uh if you do it. And so uh this has been like quite a theme for me. So I guess I'm in my 10th year of

trying to make PyTorch models fast. In the beginning the magic was compile it with the right libraries from uh uh the GPU vendor and then a bit later I tried to learn CUDA and write some kernels. Chances are uh that you've used batchorm or CTC loss and if you were angry with that it's probably my fault. And then later we did torch script and the LSTM fusion

things. Um and there already we tried to like automate it and then with PyTorch 2.0 uh we had torch compile which did most of the heavy lifting. uh with our friends at NVIDIA, we uh built Lightning Thunder uh as a thing that optimizes models with Envy Fuser and Transformer Engine and other libraries. Um and so today what I'm trying to do is take some lessons from lightning

thunder and use them with all the magic that torch compile offers. Um and I should say like I put this as a journey because PyTorch conference for me is also always meeting uh all the generous people I learned from over these years. Uh and so I'm very very humbled to be able to be here and out of pure luck and generosity uh of all the people that

mentored me. Thanks. Okay. So we have the source code. we want the execution uh and that's how we did it in 2017 uh and now what we do is we have this compile step um so in an ideal world what I could do is write the model like I would in math like basically copy the the paper turn it into some NN module things um and then

I would scale and optimize for the hardware and so what I want to do is I want to decouple this part of writing the model as it appears in the paper um and then uh running it And how to compile solves this problem is that you compile it and then you get this optimize module and it will call uh into a Dynamo a front end that runs

the Python and observes it produces a graph model module sends it to an a backend typically inductor out of the box and then it ex executes an optimized function. However, like this helps a lot, but there is a problem with it. And so to show you what the problem is, I randomly looked at some of the issues that were tagged with road map uh in SG lang

that I happen to have open as a tab. And so when you look at it, you have lots of issues uh where it says in some form of another like we have this optimization but it only works on one model or with one quantization or with one sort of distributed. And so what we need is we need these composible optimizations that allow us to write this optimization

once and then run it on whatever model in whatever configuration we need. Um and so this is like SG lang but I could probably pick any any larger uh framework vlm will have it too. Uh and so this is one of these problems that you have uh uh when you have too many variants of the same thing. Okay. So if we want to talk about like how

do how can we compose how can we apply several optimizations at the same time um we can look at what the NN module does and to my mind NN module does two things. It keeps the parameters buffers along and then it has the forward that tells us how to how to run the code what to compute and so parameters and buffers have very limited properties at least

if you don't do fancy subclasses like Ed does did in his keynote. Um but so basically when you have those tensors you have shape device and d type and the fancy people also have uh stripes but uh that's about it. Um and so changing tensors and adding new tensors to a model is fairly easy and straightforward. Like if I want to shard it for FSTP, I need

to like throw away the parts that are not going to live on my device and keep one bit of it. But the raw PyTorch code in NN modules is much more difficult to compose because we can only replace but we can't change the code. And so we have this end time M problem for code. When we use torch compile, we solve this because we get the graph

module where we can easily edit the code. But what should the initial trace be right? Should it like use the transform model then it would not run properly because uh it's missing the information how to handle the transformed parameters or I can phrase it in an other way. I can change it but where do get do do I get the inputs from and we cannot easily change

the inputs after compiling the model. So in essence I would be a lot further if I could use Dynamo on the original model and then inductor with the change parameters and buffers and obviously like we've heard that it's moving in that direction uh uh this torch compile as a toolkit and so here's like what I want to do with this toolkit when it happens and I hope

that I can convince people that it's worthwhile enough to to support this as a use case. And so if we look closer what the output is of torch compile, it's not just the graph module, but it's also guards like the guarding the cache entry to decide whether to retrace, recompile or not. And it's also a wrapper that basically replaces the entry point in PyTorch speak in the

bite code uh with a thing that grabs all the parameters and then calls whatever the back end produced from the functional graph module. And so we have these three bits that we somehow need to take care of. If we can do that then we win because basically we can go and we take a compiled module class that basically takes the original module runs Dynamo on it uses

the original model also as like our base transform module and we have guards wrappers and the G graph module and now we can change all of it the transform module the guards the wrapper and the GA graph module and we can then send the result into the back end into inductor for example. And if we do this right and we have a contract like that whatever comes

out of the transform is also something I can stick into the transform. We can combine these easily and this is like what we want to do. And so to make it more concrete let's look at FSTP and materialization. So say I have a model that doesn't fit on one device and I want to run FSTP on it. What I need to do is I need to short

the model and then I only need to materialize the part that lives on the device. And if I can do that uh I'm good. And so I need to also alter the forward graph to insert the right all gather operations. And I need to alter also the backward graph to insert the right operations to get the parameters for the backward if I need them from safe tensors

and also to uh scatter the gradients to where they belong. Um, and then if I'm at like systematically transforming, I can relatively easily also handle the state dict because I now have a state dict for the transformed module and I have one for the original module. But if I teach the transforms to take one step on it, I can also compose it. And so I can load

a checkpoint. There's basically just one thing for the entire module and the loading will automatically take care of this of the shorting or when I save it I can automatically take care of the unshorting. Okay. So, one of the things that I struggled with and probably for much of the reasons that ad could uh uh explain much more in-depth is that to my mind like solving this

with uh tensor subasses is a bit tricky because the decomposition needs you to iteratively unwrap and basically introduces a fairly uh uh complex handling step there. Um there is tensor flatten it and helps but uh sometimes there's gaps like for example tensor flatten for n parameters doesn't even exist. Um and so this is something I looked at and it didn't work immediately so I skipped it. probably

someone who is smarter than me could make it work. But so the idea is if we have these composable transforms, we can do all sorts of things. We could do autodiff in there. Um but we can also delegate that to AOT autograd. We do mix precision low precision quantization offloading to CPU or uh uh checkpointing recomputing for the backward. we can do distributed uh and so on

uh in the in the transforms. Okay. So, but talk is cheap and code rules and so if I want to do that uh I know what to do with the graph module what do I do with a guards and wrapper. Uh my idea was to just ask uh on the PyTorch dev forums. Uh but it sounded so weird that uh basically people didn't know how to help

me. And so what if you if you look at how this works internally there are quite a few moving parts. So you you have PyTorch the Python code, you have the Python interpreter frame evaluation API which is uh something that CPython produces and that Dynamo uses to switch uh the code objects underneath uh Python uh in order to call into the compile module instead of the regular

one if you want. And so what happens if you call module.forward it goes into CPython then to Dynamo C++ it checks all the guards in C++ ideally if if that is possible. Um and if it if the cache entry has a cache hit we're good. we can just call the wrapper code and if it doesn't it will call Dynamo tracing which produces guards and wrapper code uh

for the new cache entry and then it also calls the back end for example inductor to produce an optimized function that then the wrapper code can call. So here we have what we want to modify um and we can easily modify this part that goes through the back end because we can have a custom back end but we don't know really how to modify wrapper and guards

or I didn't find it really and uh so the other thing I want to do is I want to switch I want to have the original module and tracing for the for the left side and I want to have a different the transform module when I call the wrapper code. And so after for preparing this talk I like found a hugely complicated way to do this. it

had like 500 lines of code just for that and I was very very unhappy so uh even though I didn't find an official way to do it uh uh I found a function called convert frame and apparently it had all the things that I could that I needed. So uh if you replace that monkey patch python uh or monkey patch pietorch here uh you can hook into

it and you can have a transform a change of the of the way it works uh that is like not super safe uh but it gets things done. So we can just replace the closure in the object that gets called in order to switch the model to the transform module. Then there were some guards failing and I just deleted them. Ideally I would learn how to edit

them instead. Um for the forward I can just call the torch compiled module does all the things under the hood. Um, and then I have a custom backend that basically applies one transformation after the other. And this will do the code transformation that we saw, which is kind of the easy part. But, uh, if you transform the example inputs with it, which you consistently can do, then

you can send the result into inductor uh, and get the nicely working fully fast uh, function from it. Okay. So what do I get out of it? Uh what I do get is I can instantiate a module and if it doesn't fit on the device, I can just instantiate it in fake mode. Uh and then run it through um my compiled module with given transforms. And so

here I use two transforms. the FSTP to distribute things and then a materialization transform that also hooks into this mechanism and it runs the init module by module on the original uh uh module as it is written by uh uh transformers or whatever library or the programmer and then it will push it through all the uh state dick transformations in order to initialize uh the transform module

from it. And basically I can just call it just like I would call any uh transform module. And so the transforms are simple enough. All the bits that I want to change I can define functions that change this bit. So I have graph transform but also modules loading state dict and for saving the state dict I can also have a reverse transform and to see it like

it's very simple but in the FSTP for example uh what you do is you change the code and all it does it inserts some sync points and the good news here is that uh the FX graphs have been designed to be changed and so there PyTorch has really really great tooling. Uh and that might even not be the most elegant way. Um and so that's how the

code works. And so what you get when you're using this code and hopefully when we can have a more official version of uh uh using this code is you have this model times quantization times optimization times distributed problem and if you can write them as composition of independent imple transforms uh that are separate and they have a have a good before and after state. Uh you can

basically mix and match them as you want. And to my mind like model changing transforms as I describe them uh provide this very elegantly. And then of course this was like a demo that I chatted with Francisco and he knows a lot how to make this uh much more robust. Um and there's a a repository auto parallel from Meta uh in Meta PyTorch that uh does similar

things. Um I would need to learn about how to do the backward and uh hopefully I can pick his brain for a while there. Um, and then there's uh you could also do things like multiple backends for custom ops uh and port more of the transforms that uh some of the great colleagues there here too and from Nvidia and us did for Thunder and you would get

something that's very similar to Thunder and Spirit. Um, one of the other things here, at least for me personally, like uh uh the ability to uh uh like dive into things uh interactively from a Python Jupiter notebook interface uh uh has been something that I like a lot more for exploring these things than uh than just dumping things. uh but yeah there is a question of style

too and that's why I wanted to to say about the uh uh composable transforms hopefully I you found that it might be a good idea and a idea worth uh hanging on to and maybe at the next PyTorch conference we can have like a safe implementation of the same ideas. Thanks so far.