Bay.Area.AI: torch.compile and vLLM, Antoni Viros Martin

Bay.Area.AI: torch.compile and vLLM, Antoni Viros Martin

Recording: Bay.Area.AI: torch.compile and vLLM, Antoni Viros Martin

I'm a research scientist at IBM um based on New York City uh well technically not New York City it's Yorktown height which is like like 60 70 miles away from it but you know pretty close to New York City in any case um what is my background so I did a double major back in the day in computer science and aerospace engineering I then uh moved to the US for a PhD in Aerospace engine actually um and you're like like why are you like not making rockets and spaceex instead like why are you here right um so the reason I'm here is my PhD was actually based on um building chatbots for engineers building satellites right so I did a lot of work uh with not a lot of resources because when you try to like ask uh Aerospace lab to buy you a GPU they're like why do you need a GPU you're you're making like not like you KN three Graphics or anything like you don't need a GPU right so I had to make do with what I had and I had to uh optimize models a lot to make sure that they were running well with um you know very not good hardw right and that landed me a job at meta eventually and with the layoffs uh they hit me as well and then uh uh I ended up at IBM uh doing basically the same work I was doing at meta which is working on Pythor and optimization for models and that's why I'm here today uh just going to be talking about how to optimize models specifically for BLM uh so right let's get down to the talk what is the objective here today so the first thing you need to know is well we want to squeeze the most performance out of a model right like we we turn a model it's okay I'm just going to turn down the M okay no okay um can you still hear me okay sure all right um so we want to squeeze the most performance out of uh each for PR in a model right um but first of all we need to Define what performances right in in a case of a language model usually what we measure as performance is two different things one is the latency which is how fast you can generate a token so every time you get like a prompt and you want the next token in that prompt which is how a language models work you want to know how fast you can do that right the other thing is throughput um usually if you're a company that you're that is selling inference for uh for any kind of language model you want to use your hardare to the maximum right and that using means uh maximizing the throughput which is how many tokens Can you get out of a single machine for a certain amount of time right usually you measure like the tokens per second um these two things are usually like they are like you know like better latency can mean like better throughput but that's not always true sometimes you can get like a very good uh C optimization that will work for generating a single token but will not scale well to like generating 32 tokens at the time or 64 tokens at a time right which would make your throughput much better um and that means that you know you have to like play a game of like do I use this optimization or like what look does my work look like so that you know can I get like everything running with a maximum latency and through put at the same time um of course there's like Myriad ways of optimizing models uh and today I'm just going to be talking about one which is the P compiler uh being used to reduce the latency of BL custom C kernels so what is torch compile or a torch compiler um I think it was about two years ago and if someone is here from meta please correct me if I'm wrong um but it introduced ERS 2.0 the first table version uh torch compile is basically a compiler for machine learning models so you can get a machine learning model like you know defin in Python like you can see in the hopefully you can see my pointer yes so defined here you have model in Python and then the idea is you want to convert that into something that uh some kind of AI accelerator can run like it can be a GPU from Nvidia it can be a GPU from AMD it can be a CPU from Intel it can be like any of like a inferential it doesn't matter right um and and the idea is because the going from python to a machine code is not a very straightforward process they do the same thing as most other compilers do they Define a set of intermediate representations each one closer to the hardware and then they find a set of translations from like you know python to the first IR which is the which is known as the FX gra and then from the FX gra to the next IR which is the a or the prim AR in P in this case and then there's another like lowering from the this uh a IR prime prime IR down to the hardware which is you know either Tron kernels or like ca kernels or like depending on the hardware you're running right and that would be like what's known as torge inductor for like Tron in this case um the good thing is that you can as it was defined uh and designed in by torch it was designed for both training and inference because before torch compile there were other compilers out there the issues that most of them were only for inference uh this can also use for training although in this case of VM of course that's not the case that's we're not worried about that um and of course the main advantages are you know it's open source because it's part of pytorch uh they have pretty concise apis uh it is flexible in the sense that it supports things that back in like when this came out at the beginning other comp didn't have like support for dynamic shapes like you know you can grow your your tensors and this comp will adapt you can compile a model that sometimes fails to compile and then the compiler will just break it in different sub models and run like each sub model at a time before like between the brakes and everything which other compar would just like be like no I'm not I'm not doing that and it's pretty easy to add new hardware all right so what is pitch attention um they asked me to introduce this a little bit because you know I'm the the first Speaker today so this is the K we're trying to optimize right um and this is moving a little too fast to actually explain what's going on probably oh but but the idea is um when you usually do a run a Transformer model uh there's an operation in there that's called attention and the idea is you get like uh three things uh the query the key and the value and you you need to multiply the query in the key and then apply a soft softmax operation on that and then multiply the result of that operation times the value which is like you know two two mat multiplications in a sop Max operation pretty much right now um most attention like uh computation and before page attention required one thing which is that the query the value and the key were tensors in the GPU memory or whatever you were running those that and they were contiguous so you had to have every single token in a sequence together that's fine when you know you have like 256 tokens or like like 512 tokens when you have like 40,000 tokens on your query or like your kill on your value um continous space of that size in the memory starts to get a little scars so that's when something like pitch attention which uh allows for this computation attention to be done uh with blogs of memory that are not continuous in memory uh is where like you know you can get an advantage here because you know you can just use the a small spot in memory that you can just like um take from and then the vention k which is written in CA in this case will take care of like multiplying the right places in memory for you and getting making sure that the attention you get at the end is the same result or similar enough to the other attention kernels like let's say F attention or like the original attention implementation for example and VM FS if I I messed up here please correct me but yeah um they say it's [Music] okay all right but yeah that that's basically what we're trying to do so the way this currently runs is uh you have your py code and then you at some point you call a function that's called page attention V1 or page attention V2 and then this something similar to this gets run on the GPU right um so when you try to incorporate that in torch compile the torch compiler needs to in some way understand what this means to mix it together with every single operation you're every single operation you're doing in the in the GPU or any like you can have like Matrix multiplications before like uh tensor views or t on shaping or like you know concatenations and the Dodge compiler will know how knows how to work with all these things right what it doesn't know is how to work with pit attention and there was an attempt to actually make VM uh care notes like pitch attention be understood by the to compiler uh but there were like some issues right um the first issue is that the py compiler is actually just in time what that means is that you cannot like like you know in C++ you you have your code you compile it and you get an executable and then you can run the executable forever right that's not the case with the py compiler in the py compiler every time you start VM or like any other software that compiles a model you have to compile the model again and this adds like a significant amount of time to stting them right because compation is not a free process so in the case of VM uh loading a single model would take like between 2.5 to 6 minutes of like you know just from when you start running the pr to like when VM starts running and being able to serve um tokens that seems fine but like right now probably LM of time is around like seconds and not minutes right so this is a significant amount uh second issue is that um the pytorch as a library only has operations that get inputs and then speed outputs without modifying the inputs uh they shows that BLM kernels in general actually are in place which means that they get inputs they don't return anything but they modify the inputs in place uh so there's a you know the to comper doesn't know how to understand that or didn't know how to do that and then finally the torch uh P toor custom Ops API had issues on like how to like you know there's a lot of overhead when when you try to launch this custom kernel with the py JPI to do it so that the compiler understands this adds a lot of overhead that's uh unreasonable especially for something like BLM which is very uh particular about like having the least amount of latency possible for the model right all right so point one is still not soled today uh there is something called Tor export that some people are working on at meta uh that should solve this point1 which basically to compile the model and get executable and then just run that executable but until that is stable and uh easy enough to use and works well with all these models we're just going to focus on the second and third points so how do you launch Uh custom Kel efficiently with torch compile um so the way you usually works is like torch compile is actually capable of like you know when you have a function in Python that calls another function that gos another function it's capable of like removing all this C stack uh overhead and just creating a single function or like a piece of code that can run all your model without any overhead right uh the issue is that you need to teach uh torch compile how to do that with your operations because he knows how to do that with the p torch operations not any other one so then you need to do three things and again if there's anyone from meta here that worked in this do feel free to correct me if I'm saying something wrong here uh but thing number one is it needs to you need to tell the torch Library how to what inputs and outputs your function has you need to two um tell p how the outputs depend on the inputs in terms of shape type etc etc and number three you need to change your function to actually return something because if you don't return something by will consider that function to be that code and then just the compiler will remove it from your code pretty much because if you're not doing anything why is the point of calling something right because again it assumes that everything is returning something by default so this is the code for step one uh you you just Define like this Library this tor Library object with whatever code you want which is in this case page Detention of course and then you define a function I'm using another function that VM has which is called rehap and cache which is the one that stores your new keys and values to the KB cache uh the page K cach with like this like um non-continuous storage blocks basically right um I'm using this one because it's a very simple function that is easy to show in like slides uh because the pension kernel has like twice as many parameters like a lot more code and everything so it wouldn't fit as well here but the idea is you define the function like this like you know it's a ship and cache it has like five different inputs there are all tensors and then it returns a pair of tensors which in this case are the the the key cach and the value cache and then on step two uh which is the telling it how the input outputs depend on the inputs which is saying okay the key cach and the value cach are basically the same as the inputs as far as like ship and everything is concerned but you need to make sure they are contigous in the memory which is the how uh currently VM defines it right and then last but not least you need to tell pyo how to call the VM kernel that is in place and turn it into an out of place kernel like you know something that returns some value right so it has like this last part here step three in which you uh Define you know ensure everything is contigous just in case because sometimes uh the torch compiler likes optimizing memory layouts and sometimes that means making continuous memory noncontiguous if the operation you're calling can accept it in this case P attention not accept non-contiguous inputs so you need to make sure they are continuous then you call the original function which is this cash Ops sh and cash is a C++ qua kernel that this gu developed and then you just return the key value in the key the key and the cache value the key and the value cach sorry now um that's not it though um so remember how I mentioned like there's like a bunch of uh intermediate steps in between the python code and the Machine code right uh what I just did just covers the first two three steps basically but there's still the step of like going from um the torch iten uh API basically intermed presentation to the actual tyon kernels or like uh inductor code generally code that gets actually run on the GPU pretty much right uh the way you do that is you have to register something called a lowering uh which is what inductor understands as like you know this is what needs to happen when I compile this function as part of everything else how do I optimize this function right uh in the of custom C operations like fash attention uh or like you know page attention or the shape and cash things that basically cannot be expressed in terms of other Pythor operations um you just call you say okay I'm I'm going to fall back and just trust it as a blackbox pretty much right so this is the way to say uh inductor please treat my function as a blackbox just call it I'm going to tell you what interface it has and what things are important about it but other than that just trust whatever CA C++ Cod is running to just be the most most efficient way to do this uh so the way you do it is uh and I'll go to what to what this P attention Kel object here in Python looks like in a second but basically you you just say okay I'm I'm calling this page attention to shap and cache function with all these parameters and I am mutating both the key and the value cache so like these two things are going to be the memory going to be modified so please do not use that memory for anything else because you know sometimes you can py will try to also optimize that out like you know if you can use memory for something else while this is running because you already copied it somewhere else it will do it um so this is the P tension kernel object so as you will see uh it inherits from something called the fullback kernel which is basically the way of again telling inductor that this is a blackbox and it should treat it as a blackbox and just trust me to run this the best way possible right um You also see lot of code here a lot of a lot of this code I just copied from the only other operation in py that is also in place so I I I when I said like py allation and py are like you know something that get some inputs and speed some outputs I like there is one exception to that role and it's the distributed operations so let's let you know when when you when you have like two gpus and you want to add together the results from the same tensor and two different gpus you call something that's known as old red use right uh the old reduse kernel is the only operation in all P that is also in place so of course the P developers when they add support for all ruse and like all these other all scatter redu and all these other operations that are like multi-gpu operations in py they had to add a bunch of like uh tooling around them that now I'm taking advantage of to make all these in place SCS in BLM work as well like the same way they use like for use and everything else um the only issues that uh these two things I'm using here the non layout here and the mutation output here these two things are not not available in P 2.1 good news is p 2.2 just released yesterday today and those two things are available in P 2.2 so you don't need to run on N anymore to actually take advantage of this uh particular feature um so let's say you you on your torch compile uh you have your model and this is not a VM model actually this is a model from IBM that we' been running to test all the P attention Kel code before making the pr to BLM and contributing the code back but the idea is the does all these thing I was just talking about work right and then as you can see here here uh you see some Tron code Tron canel being run on top which is uh in this case I think it's the rope and medings in case you're familiar with that what that is um then we're calling the reip and cash uh which you know there's no like extra stuff going on here other than that then we're calling the pit attention a kernel itself and then finally we're calling a matx multiplication which is also like compiled by you know uh the the titon compiler to make it hopefully uh efficient although we had this discussion this morning in which we said it's probably not the best show it's Tron is not uh like uh the inductor compar is not the best choosing the best Matrix multiplications operations yet but what does the performance of this look like right um well the test I was running is on a bat size of four and I generated not a big batch like nor only 356 tokens um to begin with uh and this is running on an A 1 a 100 GPU and py 2.1 still um the actual machine is the AWS p4d in case you're fam with that it's the one machine that have a that has like 8 8100s and I think it's like um I L CPUs from intel if I'm not wrong uh and like one tab of Ram or something like that um so what does compilation time look like in this case it's again as I mentioned the the startup time is still pretty bad it's like two minutes just to compile this one model um the iteration time um from the eager execution to the compiled execution for a single token generation goes from 96 Ms per token down to 30 MS per token and I forgot but this is a Lama 7B model um so that would be like uh you know in this machine particular for this particular model that I I have it is open source so you can check it out and I'll have the link in the last slide but basically this is a 2.5x speed up um if you only look at you know this is from the beginning to the end of the for not just the for pass like it this is like I finish getting a token and I start getting a new token this is like the fulltime right if you only look at the part that we're compiling which is just the forward pass for the fun for the model the that time goes from 94 milliseconds to 30 milliseconds which is even better like the issues that you know there's some other like uh limitations to the model that of course show more when like you reduce this overhead right uh I did look at the memory requirements they basically stay exactly the same down to the bite so in this case uh the toor comp didn't manage to optimize any memory out uh but it didn't add any extra memory on top which is good it means like you know uh we're not like uh losing throughput at all because usually throughput is memory memory bandwidth limited right and then um the main issue is that as you see in this picture this is like what the GPU like the bottom line is the GPU execution what's being run in the GPU and the top is what's being run in the CPU as you can see the CPU is pretty busy all the time as you can see the GPU is not that busy so uh there's still further ways to like reduce CPU overhead uh the VM folks will talk about it but they they went the other way around like there's the the next step to do this is using CAG graphs uh to further reduce the uh overhead from like all these CPU operations and it to zero basically because you only run one launch for a Kel instead of like launching every single thing from the CPU um but you can do it the other way right we can have CAG graphs first and then add compile to it in in our case we started with compile and then we next step is doing CAG graphs I just just didn't have time to get it working before coming here today but ideally when you run with Cuda graphs you will see that the the blank space in between like each of these small Blobs of actual competition in the on the GPU will get uh removed uh some other numbers that are not here in the slide but uh we checked I checked this morning is um the eager GPU time uh so the work they doing in GPU on eager mode for pytorch uh is 12 milliseconds um out of those like 94 on compile it's 9 milliseconds out of the 30 so there is still like you know compile is is adding some optimization on the GPU as well even though like most of the optimization is coming from the CPU side uh that being said and because I'm running out of time um the conclusion is basically yes the compiler can offer significant gain in performance although it should be probably combined with CAG graphs especially for like smaller workloads like this um but it also is still very much in development and you know these numbers will probably look different next month with like the night L build like you know by 2.3 when it is in like uh May or so it will look very different from what I'm showing here I'm sure like probably better because compar keeps getting better um but yes there's some issues and then of course the BLM got cut out on some of these issues and then thanks to them getting cut up on that we had like a case to like make to like some P developers to improve some of these things and that's why we can it working now right and um why would you want torch compile well uh Tor compile as I said is very good because it lets compile to different like backends and that means that you can probably get a lot of uh support for like AMD or Intel or other CPUs and accelerators for a lot cheaper than it would be like rewrite every single custom Kel that you have in your code for each of these different processors like uh VM does have a lot of custom kernels much of them most of them are Prett in Cuda right now and probably if compile can know for a similar level of performance but also on like five different like accelerator that's something that BLM may be interested in and everyone else here in the room should be also interested in especially if you're part of the AI Alliance um you can check out this code uh you can use the the QR code and the link are the same um so use whatever you want and then uh this particular work I just presented is in the this Branch uhe to a picture so you know which branch to look at there also PR uh open right now which is called like page attention or something which points exactly the same branch so yeah I think that's it pretty much thank you very much and I'm going to leave the FL to the B folks some question um so while we set up sure if you have any question for Anthony I'm here it's a good time to ask if not after that you yeah I'll be here until 8 to so oh uh so the question is uh which operations are being fused that actually get dis speed up in in our case it's the drop in beddings which one the Rope embeddings the yeah the all the multiplications and everything that are happening on the Rope embeddings are getting fused into a single Tron kernel that runs like uh I think it was from 069 milliseconds to like 003 milliseconds so a pretty like de speed up uh there was someone else here yeah yeah yeah s my my my question is have you benchmarked the Llama 70b compilation time with for compile uh yes and I just run it like three hours ago so 70b on the night list from two days ago which is the one I think I have on my environment right now took about uh 6 minutes to compile six minutes on a p4d machine on a p4d yeah okay that's from like this afternoon so okay yeah yeah so again the question for like the recording is uh what issues have I seen with Tor export and why are we not using this for like the setup time right um so the the main issue is that a lot of the function like you know when you TR to run Lama 7v for example there was almost no support for all the colle like all Theos and all GA and all these different collectives that you have to run on um to run like any any tensor parallelism in Lama 7 for example or any other big model right uh there was no support for that I think they landed the support for that recently so I need to try again and see if that works I just haven't had the time phal time actually like go and try it again yes um I'm wondering uh how does your work handle Dynamic sizes using the Tor dynamism support or are you doing someing yeah so this question is about uh if all all this all these I just presented does support Dynamic shapes for compilation with VM and the answer is yes it does and it just works with the torch compile yeah right so for more questions later yep all right thank you very much and I'll get [Applause] you