Efficiently Serving LLMs at Scale

SBTB 2023: Nick Hill, Efficiently serving LLMs at scale.

Recording: SBTB 2023: Nick Hill, Efficiently serving LLMs at scale.

okay uh hello everyone thanks for coming um so I'm going to talk about uh efficiently serving large language models at scale um something that we're working on at IBM and focusing lots of other places too um so just give a quick introduction to myself uh I have a a cloud background I sort of work on distributed systems um and in particular infrastructure for machine learning um systems I was involved in productizing the original Jeopardy IBM Watson if anyone remembers that um and then scaling a lot of IBM's Cloud uh machine learning products and services like what's an assistant what's in Discovery um those kinds of things um I've yeah been more focused on the INF on the uh inferencing side of things um and in particular I designed a framework called Model mesh which was open sourced um and that's for for serving very large numbers of models and efficiently managing uh you know if every user has their own model and you've got lots of small ones um scaling them out and effectively kind of managing a a distributed cluster of them um and I contribute to some other open source projects if any any scar people here um I sort of committed to the netti project model mesh was was um used netti a lot was built on Java um so yeah I about a year ago or maybe year and a half ago um I switched to look at uh serving large language models that's obviously quite a big difference because um instead of large numbers of small models it's like a small number of large models um and uh it's you know very interesting space to be there's lots of um research coming out every day and opportunities to kind of look at how you can apply engineering to the to the latest research in the space um so in this talk I'm going to sort of give a a little bit of a journey um of what we did to optimize like llm inferencing when you know we first got involved um as part of that explain one of the the sort of key techniques called continuous batching um and then I'm going to sort of Depending on time I'll go go through a bunch of other techniques that um have emerged since since then as well um so if you rewind to Summer of 2022 um this is before chat gbt came out in November December of that year um there open AI had a had gpt3 but it was not that many people knew about it it was it mostly a sort of research thing um but clearly like a lot of people were seeing the potential um but it was uh it sort of came with restrictions um in terms of that what you could do with it commercially so uh there was a a large open source Effect called Bloom that some of you may be familiar with um it was a open collaboration between lots of organizations and they they trained 176 billion parameter model um that was meant to be uh sort of the open source answer to gpt3 um so you know this this had been trained but uh because it was so large no one you know there was not much um known about how best to to deploy such a large model um the weights were like 350 gabyt or so and if you want to serve it in gpus there's no gpus that large so you have to kind of have a whole node full of 80 gigabyte A1 100s and and spread all the weights across um these and in IBM we wanted to give researchers access to this um obviously they can't all deploy their own because it the gpus cost so much and and are so scarce um so we were looking at okay how do we sort of deploy some like a Bloom as a service um so so uh initially um there was a implementation of Bloom in the hugging face Transformers Library this is what most people used then least and still uh for the most part for sort of um working with Transformer models uh but the size is a problem so there's sort of these other Frameworks that you can be used with it um one is called um accelerate so that allows you to split up the layers of the model across multiple gpus uh the trouble with that is it's kind of slow because it does one layer at a time so if you're trying to do real- time inferencing it's not a great solution uh the other way is if you um split up the model uh by layer so you split each layer um across you know each of the layers is split across the GPU each GPU then you can run the layers in parallel um or you can run each layer in parallel sorry uh and there was an implement this is called tensor parallel there was an implementation of this from Microsoft um called Deep speed inference uh so this is something we started and we sort of wrapped it at a web server um and had a little queue in front and users could sort of submit um batches of of uh requests to to um be influenced by the llm um and we passed this to the hugging face uh generate method in Transformers um this you know this first implementation was wasn't not ideal because once you have more than one user um things things slow down quickly uh it takes a long time to do this alter aggressive intering which means you it it's being done token by token um and so for you know if you're generating a large amount of tokens that's going to take a long time and if another request comes in in the meantime it has to sort of basically queue up and wait till that's finished um so this would result in a you know low GP utilization uh low throughput um most request a kind of single input or like a low batch size um so uh instead like what a lot of the servers do is sort of dynamic batching where you can kind of Auto batch request requests that come in put them into a batch and then like pass that through um the Transformers Library took you know batched input um this is okay if all the requests are kind of similar but if they're different lengths then it's it's kind of a problem because the the short ones are going to wait for the longer ones you know the the the amount of time it takes to run a request is basically proportional to how many tokens the request generates um and also the the parameters like if if you any of you guys have uh used um uh chat gbt or well actually not chat gbt but um some of the uh other hosted llms uh there's parameters like um the temperature and top K and top p and things like that and if the requests have different parameters um they can't be combined into the same batch using the generate method um so solution to this was to split open the generate method and inside is there's basically a loop where you're generating each token one at a time um and in and instead of applying the parameters to choose the next token um in the same way across all of the entries in the batch uh we sort of keep track of the different parameters and we can do those one by one for for each of the elements in the batch um this is also nice because we can then stream the output now we're in we have uh uh you know in the code that's actually doing this Loop we and we can then also uh finish um the shorter sequences as soon as they finish rather than um you having to wait for the longer one so that was a a a big Improvement um there's still some problems here because if you have uh um uh you know very different input and output links then the batch is still going to wait for the longest one to finish before you can start a new batch um so you know some would finish a lot earlier than others and then there's there's requests waiting in the queue uh so um the idea is to uh to sort of key idea is to modify the batch mid sequence um so you know as sequences in the batch finish you can prune them out and then add other add other um sequences into the batch uh which sounds kind of simple but actually it's there's a little bit more to it than that um so uh one of the key optimizations that's used for inferencing is called KV caching um each uh like every time you generate a token you have to you B basically calculate the attention that that to uh the attention between that token and all the preceding ones in the sequence so there's a lot of redundancy um and you can cache like everything up to the current token uh you can cach all of those weights um so typically you you know you you process all the tokens as a first pass that's called a prefill and then from then on in you just send in the next to next set of tokens for each of the sequences in the batch and you you keep this cache of tensors that have have all these weights um so the trouble is if you if you're sort of processing one batch and then you want to add some more requests in you need to kind of first process the get the the new batch into a similar State uh as the old one um so the way it works is you you you you run the prefill step for the new batch and then you can basic do some sort of tensor manipulation to concatenate the KV cache of the the new batch with the existing batch um to to sort of form a new combined batch and then kind of continue from there with the the token by token can decoding um this is kind of a visualization of that so um at the at the the top you've got a batch uh the the zeros are kind of represent padding tokens because we you know pad the the shorter sequences and then the sort of orange column uh is the new tokens that have been chosen for each of the sequences in the batch um and then you know you see that one of the sequen is complete uh at that point we can kind of prune it out of that first batch and it means we have a bit more space now in memory cu the the problem with the cach is that it it takes up a lot of GPU memory so that's that's typically the the main constraint for how large batches you can run which in turn is the constraint for the amount of throughput um and so now we've now we've reduced the memory that the KB cach takes we can uh pre-fill a new batch and then and then join the tensors to create the the KV cache um which is the at the bottom there uh and then continue so what you have in effect is kind of a never-ending batch um where you know as new requests come in you sort of add them in if there's space and then um as they complete you prune them out and then they just sort of all into leave so you never really finish it's more like a train than than a batch um and uh yeah a lot of the optimizations are sort of um centered around how to best deal with the KB cache in a way because manipulating the cash is is it's because it's a large amount of data in the GPU that's one of the main overheads um so going back to our sort of Bloom as a Serv uh we were sort of looking around at the time for um for how we could improve this and there hugging phas came out with a Blog uh of how they had done this they sort of made their own tensor parallel implementation um and they had um um they' done a lot of the the things that I'd mentioned in terms of splitting open the generate method they hadn't so that you know requests that are finish could could finish early and also they could you know have heterog genus requests with different decoding parameters um so we we B rebased our code on this um and then started to look at okay how can we do this technique of like adding new new uh sequences into the batch that I described um around that time um a new project sort of appeared also from hugging face called text generation inference um this was like a a redesigned sort of production grade server based on the same thing um and that had a front written in Rust which was decoupled from the backend um sort of python Shard processes that actually uh do the tensor manipulation and and run the the model um the front end sort of dealt with the batching and and the the request queuing and that kind of thing they'd also Incorporated the the uh this sort of never-ending batching logic um so we quickly adopted this uh we um actually forked it ourselves and you know started to make some of our own improvements on there um so we sort of now work work from our something that was orig originated from from this this project um but we also sort of work with them originally um and contributed uh a number of things back um so the term you know we that was adopted for this I think I think I was the the first person to kind of think of that name but it got picked up and it's it's used now and lot of the other um uh projects that are doing a similar thing uh is continuous batching um there was I think the earliest reference is actually in a paper um they didn't use the same name but of a technique that was similar to this um called iteration level scheduling uh it was a paper of uh called Orca it's like a not an open source project but it's it's a yeah a paper that was like a came out a few months before that um so the next thing we looked at was um how to how do you decide when to add uh new sequences into the cache so you have a cue of um of these new requests come in of sequences of different lengths and they all have different max number of tokens that they're going to be generating uh how do you how do you figure out how how many of these you can add um and typically when you the servers do Dynamic batching you just set a maximum batch size and you just sort of fill up the number of sequences until you get to that size and as more removed you can add more in um the the trouble with that is that uh the memory that is the limiting factor is a function of the sequence length as well as the batch size the number of things in the batch um so you know if you were restricting the batch size you'd have to figure out like okay what's the longest sequence length uh I I want to support and then you need to size your batch quite small to accommodate that um so one of the things we did next was to introduce Dynamic batch sizing um which decides whether or not to add sequences into the batch based on uh how many input tokens they have and how many output tokens and you know they're predicted to to produce and so we can figure out okay if we add this would would we exceed the the size of the KB cache as the sequence is generated um and if not then we can add it in otherwise it sort of stays in the queue so the nice thing about this is we can support much longer sequence lengths uh and the batch size will just sort of flex um to accommodate that if the the sequence lengths are very large then we'll end up with small batches but if they're sort of a a medium size or a small size then um the batch will kind of flex up uh and we can um we can have much greater throughput uh with with much larger batch size is up to say 128 or something or even 256 um okay one of the next uh optimizations was uh and this was one we um that uh came from the the Upstream sort of hugging face project when we incorporated it um but uh flash attention is say um is a a technique that really optimizes how the Matrix multiplications are done uh in a Transformer uh how the attentions applied and sort of making use of the the layered um GPU memory caches um and there was a a good implementation of that I mean a sort of reference implementation the um uh which allows you to provide the cache in a kind of contiguous format so instead of having a batch Dimension you can sort of concatenate all the sequences um and provide offsets uh and so as well as as well as the actual ATT the operation to compute the attention being faster this meant that um you can you can obviously pack a lot more tokens in especially when the the the sequences are variable lengths before we'd have to sort of pad to create these rectangular batches uh and that's you know this is just showing two Dimensions but when it's the KV cache you have additional Dimensions which you know every one of these numbers these tokens is actually a big chunk of memory so you have all these wasted zeros the padding uh padding tokens um once we you switch to the sort of the variable length sequence uh single batch Dimension then we can just save a lot of memory by packing a lot more in and then you know that the the algorithm to fill up the batch is similar but it just means we can get a lot more ends that that helps a lot with throughput um another optimization is vectorizing the decoding so one of the problems when we split the generate method and don't apply the same um decoding uh logic to each of the the the rows in the batch this is where you know you're selecting from a a probab probability distribution over the the possible next tokens the way you do that might be different for each row in the batch now and so what we'd had to do is sort of a loop and uh process each one separately and each of those involves some GPU operations because these all of these scores that come out for each of them are um are intenses on the GPU uh once the batch size gets quite large this becomes a bit of an overhead um so vectorized decoding is basically where um you you create tensors that sort of combine all of the parameters across the the batch so if they all use different values for topk for example you can put those all in a tensor and apply it to the column of um uh of the the the probabilities at once so basically you know vectorizing each of those operations so that was a big speed up um so this is a just a picture of you know our version of the server showing um the rust front end uh and the the python shards uh on the right hand side um and basically the batching loop on the rust side is deciding on which which thing's the batch and like between every token it goes back and forth and and uh sort of controls what the how the how the tenses are are put together and the the forward passes are done on the python side um so something that came out very Rec recently uh from Microsoft um is a technique called dynamic split views this is in their sort of revamped uh deep speed inference or uh deep sple uh fast gen I think they call it now um so this is actually uh what you know a problem with the the the way that the continuous batching works as I described it is you have to kind of interrupt the flow of the of next tokens for the current batch to to add a new batch in and the the pre-fill operation um where you're you're Computing the the the initial weights for like a for all of the input tokens is actually quite expensive computationally so that takes much longer to generate that first token than all the subsequent tokens um and it's also Bound by the compute rather than the memory so um the later steps uh are memory bound the first is compute bound so you're actually not you don't have necessarily great um especially with smaller batch sizes you don't have great utilization for the just the next token stages um so the idea here was like to um uh sort of amortize the uh the prefill step across multiple um decode steps like next token steps so they sort of if you have a a new batch to add in you can um you can actually combine the the the the pre-fill operations with the decode ones to just enough to saturate the the compute and not slow things down um but it means that you you you're then able to soak up more of the the utilization and you have less of these kind of bubbles as they describe it um so that's like very interesting technique um for for helping with that um another one that was you know it's been a been around a few months now is a a new um inference server called VM and the the the big contribution they had was to uh apply sort of virtual memory principles to the KV cache so one of the other problems with the KV cache is um it it's a tensor in the GPU that keeps growing uh but because you when you're when you're um generating these sequencies you don't know when they're going to stop uh and so you don't know how much space you need to reserve for it and if you don't reserve space then you have to you have to copy lots of data on every single token that's generated um so you know they basically uh implemented kernels which can work with uh non-contiguous tensors so instead of having to um uh to sort of copy the whole tensor each time um you can basically divide it into blocks and uh and just reference you know chains of blocks without having to do the copying uh and then the the kernels can work on those directly and it allows for additional things like sharing reference counting and sharing the same data between sequences if you've got duplication and that that you know allows for further pack you know um efficiency in how the the KB cache is put together um another sort of technique is uh a um an observation that like a lot of the times the tokens the next tokens that you um are predicting are actually quite easy like there's like very much more obvious uh potential next tokens to other cases like if you're generating code or something like that um and so what you can do is use a small model uh to kind of predict a bunch of tokens into the future um and then use a large you know use a larger more more accurate model to sort of verify those the nice thing is you can actually because you process all the tokens in one go you can like verify a batch of of tokens usually instead of having to to wait for the next token you can do um the next four tokens for example together and and see whether they would have been chosen by the large model and if not then you can like um you know if if uh if say the in this example the in the were chosen by the large model but the the car is isn't the right one or the bus isn't the right one and you get car instead then you can just you know take take as many as you get and then revert to doing the predictions from there with the smaller model so this allows you to kind of Jump Ahead much much more quickly than than having to do these these uh steps that depend on the the prior ones um so you know since since like we got involved like uh maybe 18 months ago um there's been an explosion of activity in this space um I mentioned VM that's like one of the most prominent servers now uh and they sort of incorporate a lot of the things uh the techniques like um continuous batching and they're looking at you know incorporating all the new models and the new techniques that come out um there's uh yeah a bunch of others here like I mentioned the Deep speed fast gen Nvidia came out not too many weeks ago with uh something called tensor rtlm which is their sort of version of this um there's a project called flexflow from Stanford they were sort of focused on the speculative decoding stuff um but now incorporating a lot of these other techniques so what we sort of find is there's announcements every now and then of one of these new projects that say okay we're like this we have this new technique and we're like uh much faster than the others and they do a bunch of graphs and then the other projects will uh incorporate that technique and then maybe do something different and uh so there's there's all this sort of leap frogging and everyone trying to sort of uh do more efficient variations of the existing techniques and come up with new ones so it's kind of a very interesting space and um uh yeah it's it's continuing to change um some other techniques uh that I didn't cover uh are sort of compiling the the graphs using something like torch compile um so that you don't have to to do it in in Python like things that would take multiple GP GPU operations if it was running in Python code you can fuse into single operations um Cuda graphs is similar to that quantization is where you you know shrink the um the size of the the Precision of the the model weights and that means you can fit much more onto the gpus and so get more throughput that way uh and there's you know there's a whole bunch of uh sort of sub techniques of a lot of these that I could you know could go into um but uh yeah I'd recommend if you're interested um looking through some of these things and uh uh and getting involved cuz yeah if you're sort of engineer like me and uh it's a very interesting space to sort of apply things um and then just uh final slide to recommender this blog post came out couple of weeks ago from the Mosaic ml folks um which is a really good summary of some of the things I talked about um and sort of best practices for doing uh this kind of engineering so thank [Applause] you by