Bay.Area.AI: Efficiently serving LLMs at scale, Nick Hill

Recording: Bay.Area.AI: Efficiently serving LLMs at scale, Nick Hill

thanks everyone for coming uh I didn't do these very often and um Alexi asked me kind of last minute so apologies it's the the slides get sparer towards the end cuz I was I was trying to cram them in um and I think I didn't have time to put any any code in so sorry I'm I hope I'm not disqualified um so uh but it is technical so how many of you guys are um uh sort of know about llms and and pytorch and sort of work with those kind of things uh did any of you go to the conference yesterday the P torch conference yeah one or two okay cool um so uh yeah um I actually been IBM for um quite a long time lost lost count at least 15 years uh and most of that I was I sort of came from a cloud engineering performance uh distributed systems background and then um kind of got into the AI side of things actually at around the Jeopardy time uh when IBM Watson you guys are familiar with that W that so um was involved in some of the efforts to productize that um and then since been scaling uh machine learning Services um if you're familiar with things like what's an assistant from from IBM like I built a lot of the infrastructure that that helps um support that on the influencing side um but mostly that was like managing where there's very large number of custom models so people could sort of train their own workspaces and models um and I built a framework called Model mesh which uh basically allowed um sort of a decentralized layer for scaling out where you've got hundreds of thousands of models and deciding which ones to load um and then uh about a year ago or last summer I sort of jumped into the foundation model space before many people were talking about llms and things um and that's obviously kind of the opposite problem it's like instead of very large numbers of small models it's like a small number of very large models so that was interesting um and I haven't actually been using pi torch that long maybe any about a year since I kind of got stuck into the the um what's involved with inferencing with these these very big models um so yeah I'm going to sort of do this a little bit like a story and then dive into some of the technical stuff as I go so um yeah back last summer I don't know if any of you know about the or knew about the the blue model that came out then it was a big research effort um uh so gpt3 was around from open AI but not many this is before chat GPT um and so not many people knew about it but it it was getting a lot of you know research circles there was a lot of interest uh but because it was closed especially IBM um it was a bit of a problem for us and we wanted researchers to be able to experiment with these large models uh Bloom was kind of the open source answer to that it was a big sort of um joint effort between hugging face and Nvidia and IBM and um they trained this 160 billion parameter sort of multilanguage model um and uh the problem was once you train that how do you deploy model Al so big it was sort of kind of new territory um and uh so you know it actually because it's 176 billion parameters is that was like 350 billion uh sorry 350 gigabytes of GPU memory you needed to load on gpus so that would take um you know eight uh a eight um a100 gpus uh so you need a whole node and it's obviously not something that each researcher can can have they're even now they're extremely scarce and expensive um so we wanted to sort of set up a service so that um researchers could just get access to this a bit like they could with GPT but they they weren't really meant to use that because of the you know the the commercial restrictions and things like that um so the the hugging face Transformers library is sort of de facto library and because they had been involved in bloom um they had a sort of first modeling implementation for Bloom um but that didn't help really with how you could deploy it um so the sort of initial attempts use something called uh um hugging face accelerate which is mostly used for training like for sort of distributed training but um it also works for inference and that's where it um sort of splits the weights of the model layers across multiple gpus to to be able to fit them um that's called pipeline parallelism uh trouble with that is that it sort of runs things sequentially so you sort of process each layer in turn across the gpus which performance-wise um is very slow so a lot of this is about okay how can we um how can we deploy it in a way that's efficient and multiple people can use without having to wait sort of minutes to get to get their answers back um and there was at the time a sort of uh an early implementation in in a library called Deep speed inference from Microsoft and they had a TENS of parallel implementation um so that's something else we used we sort of wrapped these in a simple web server and that was kind of the first version of our internal Bloom as a service that that people could use um so you know we had a little queue and people batches to this thing and then they get ped to the um the hugging face generate method uh there a sort of tokenization Step um and that was kind of what we did at the beginning um so uh the trouble with this is that um uh because it's autoaggressive which means it has to generate one token at a time um it takes a really long time and if you're just queuing requests up you're sort of waiting you know you could be waiting tens of seconds per request and so if you've got lots of users it's not like you can deploy too many instances of this because the GP the model is so big um so uh yeah something else is needed um so the obvious thing would be to to do Dynamic batching this is something that other inference servers would doing at the time um or have been doing a long time but basically you you know you batch you have a queue you batch the inputs into a batch and then you can run batches and get much better GP GPU utilization that way um so that works if all the requests are kind of similar but it it kind of breaks down a bit if you've got users sending in um requests of different lengths uh because uh you might have one that's much you know much longer and then they all have to wait for the longest one to to finish um and then the way um the the generation works it's it's lots of iterations of the models for forward method and in between you you do a sort of token selection and that's the that's called the decoding step and how you select that token might vary depending on the request um so you probably if you've used chat GPT in these other models there's things called temperature or top PE or you can do like a greedy decoding all these kinds of things or you can just decide how many tokens you want to stop at or um have different stopping criteria so uh the way that the the hugging face generate method works you just you can you can um provide a batch but you only have one set of parameters so if you're batching requests from different users that have different parameters you know you're kind of stuck uh because they all have to be the same so that was that was a problem doing it you know that kind of naive way so um what you can do is uh sort of implement your own generate method so in the Transformers Library it essentially is is a loop uh where they call the forward method they get back a list of tokens like POS candidate tokens then there's a decoding step where you choose one of those tokens based on you know either some random sampling or you choose the one with the the most likely um probability uh out of the probability distributions you get back um and so by splitting this out we could kind of keep track of the decoding parameters and the the stopping criteria for each of these individually um and it also means that we can then send the requests as they complete uh so like all the requests are not waiting for the longest one um so that's a uh you know big Improvement um and then you can also do streaming output right because we're sort of in control of this Loop where we're generating each of the tokens we can decode them on the Fly and um uh and get the sort of streaming Behavior like you see if you use these llms on on the internet right um and that helps a lot with the user experience then having to wait tens of seconds to get the text back in one go you sort of get it streaming as it's generated um but this um this still has some problems uh if you have sort of heterogeneous requests um you still might end up with one very long one like all the others finish so you're only processing one at that point and your so your effective batch size is very small and the GP utilization still very small um uh and at the same time you have other requests waiting in the queue um so yeah it's not ideal from that point of view um this is sort of like a a new problem so it hadn't really been tackled before right the there's a lot of inference Ser servers at that point but they they generally were just doing sort of in a single forward pass for inference um so uh the um the idea beh behind you know how to solve this is to actually change the batch mid mid request so you have one batch that you start um and then you know as requests finish you can sort of take those out of the batch right even though others within the batch is still going still generating tokens one of them might uh produce an end of sequence token which is usually when you you stop that sequence um and the others others continue uh um and uh um yeah so and then after you've made some room in the batch essentially you can add more requests in so uh it's it's like you're concatenating the batch with with additional requests um so it sounds it sounds sort of straightforward like just add add into the batch Midway and then um but the the trouble is uh it's yeah it's not quite that simple this is where I was going to show some code but um I I'll talk through it uh so um with inferencing there's an optimization uh called KV caching um so uh when um with with causal language modeling you're always just generating the next token based on all of the tokens before so that each token attends to all of the prior tokens in the sequence um and because of that like there's a lot of redundancy in the computations like for the next token because it's it's it's uh calc Computing the attention for all the previous ones each time and those are obviously a subset each time it does it so you can sort of cach all those weights um from the uh the attention layers and then send those in next time just along with the next token so there's kind of two stages when you when you start processing a batch um you pass in all the tokens first uh from the text um and then you get the KB cach out of from processing all the input tokens and then on subsequent passes you pass that Cash Plus the next token then you pass the cash the next token um and the term there's terms for these called like the pre-fill stage and the decode stage um so to be able to add request into a batch like while it's in in progress um you need to uh it's it's kind of a problem because you you you know you've got these two things in different states you can't kind of combine the the pre-fill stage for one with the the decco stage for the other so um one idea is to like re redo the pre-fill stage recompute everything like put them all in a you you know the from the tokens that they generated so far from the first sequences you could put them all into one um one batch but then uh it's actually very computationally expensive so then you'd be like recomputing all the attention for the the um the the sequences that you already have it for in the cache um so you can so the way to get around this is to sort of combine the KV caches uh so this is kind of a representation um of that so the the first one would be like some sequences that have come in and you have a batch some are shorter than others so you can pad it on the left hand side um the sort of orange column is selecting the next token that's after one iteration for for each of those sequences um and then uh the in the the third row there's like a blue token that's like represents end of sequence and so that one's finished right so now I can return that one and the batch is now only um only two so we sort of have more room um uh and so now we can pull some requests from the que into a new batch that's that's batch B and run the prefill step to generate the cach for that one uh and then do sort of some concatenation operations on the the tensors in the cache um to com sort of create a combined KV cache so one of the the dimensions of the KV cache tensors is the sequence uh another one is the batch but then they have a couple of others for the um uh for the you know the hidden hidden dimension of the model but um if you can see we kind of trimmed off some of the padding and have have a more compact batch and then from there we can just kind of continue the process um so that means you can sort of absorb before requests are finished you can absorb more based on the space that has been created uh from others so it you end up with a sort of instead of doing one batch after the other it's just like a a never-ending batch right so it's like this train uh you can think of it and those sort of requests come in um they get added into the batch and then they can they get taken out of the batch when they're finished but in the meantime other requests might be arriving and you add them in and they're going through this you know each of them are going through the same uh loop of forward passes and where they're you know selecting their next tokens um so yeah jumping back to like our our Bloom as a service uh we um uh you know we were following what others were doing here as well like working and actually some some folks from IBM worked um on some of the early Scripts uh for um for standing up Bloom like I mentioned using uh the pipeline parallel and tensor parallel actually one of the the later talks I think is on on about tensor parallel and P torch uh so uh that should be interesting um I'm sure we go into a lot more detail on that um but at the time hugging face put together a sort of a custom imple test tense parel implementation um using torch distributed uh and they use torch. jit this is sort of before torch compile was really usable um and uh they they did this stuff splitting open the generic method so a lot of the stuff I described they kind of experimented with they didn't actually originally have the um where you could concatenate batches in flight but we based on this code and and sort of follow what they were doing um and then not long after that as we were sort of trying to think how we could do the the concatenated batch stuff um this project text generation inference came out some of you guys may have heard of that um um and this was like a better design production version of all the stuff I've described um and so we kind of quickly adopted that uh we got involved actually and had our own thought where we would we were enhancing it in different ways and started to contribute back um so we were sort of working with them uh on that um they did uh you may have heard change change the license like to a non-commercial license there a couple of months ago now which meant probably we can't really we can't continue to use it but we have our own Fork so that's you know we have a lot of the extra stuff we've added there and we open source that too now so um that's you know the basis of our sort of Bloom as a service um so we we came up with a name for this continuous batching like people were calling it Dynamic batching but it's sort of this you know unique thing for for llms um and I I couldn't find any other I sort of thought of that term and I couldn't find any references to it but it's been it's been adopted now it's like all you know all the other language servers are using this term continuous batching and so I think I'm the one that came up with it we did find a like later it's actually there's a paper called Orca that came out of Korea which um where they were doing something very similar so I think uh you know we weren't the first to implement it um or the hugging face guys weren't necessarily the first I think they're definitely the first in in the open source Library so this is the Ora thing was sort of a proprietary thing they that they described it in a paper but they hadn't actually open source the code um but now it's sort of the you know it's uh sort of indispensable like any production um llm serving you would just use continuous batching you know as a a standard plus a lot of the other things which I I'll go into a couple a couple more of now um so another aspect of this uh that you may have thought of is okay we're sort of adding stuff into the batch but I was a bit vague about when you do that um most of the time servers had like had a a fixed batch size that you configure and you just sort of fill up you know you'd fill the batch up to some maximum size and then you do your inference and then you get the next stuff and you fill it up to that size um with uh with language modeling it's like the the limiting factor is typically the memory that this KV cache takes so this KV cache that you have to sort of store pass between forward passes um it takes up a very large amount of GPU memory and that's a function of the the batch size but also the sequence length um and so basically if you have and that's usually Gates how how large batches you can have which in turn controls like how how much throughput because the up to a point the bigger the batch the higher through higher the throughput um so uh solution to this is to is basically decide whether to whether or not to add things into a batch based on the sequence lengths of those requests that are that are arriving um so uh um yeah so you we we sort of Define a maximum weight for the batch that corresponds to the amount of GPU memory um and then as as requests come in uh we look okay well how many input tokens how many output tokens are they going to produce and then kind of do a projection to see okay if we added this in now would the batch and like we know when we're going to be taking them out in future would the batch ever kind of grow beyond the limit basically and if so then we wait and you know just wait keep going and then as as requests finish then we can add them in so we can kind of keep the batch full and that the um the actual batch size I.E the number of requests in the batch sequences will kind of flex based on the the sequence length so you could have um you know short a bunch of requests with short sequences that would mean you could have a very large batch size because you're sort of multiplying them together right and so you get very high throughput but at the same time we can allow requests that have a very long context so if they sort of you know pass along document then you want to do some generation based on that or generate a very large number of tokens we can still support that the throughput at those points would kind of go down but you know it means you can you don't have to have a sort of lowest common denominator so that's something I added to the you know this this TGI project and that you know was important as well for the the throughput um another thing which was added uh was sort of incorporating um uh flash attention which is a sort of much more efficient way of of computing the attention and in particular it the library um has allows you to have a kind of uh KV cache with a um a single batch Dimension so you can just pass offsets instead of having to have uh a rectangular sort of tensil like this so the trouble here with if you have variable length sequences um the C you know these I'm showing kind of two Dimensions here right the sequence and the batch but then there's other dimensions that are very large and all these zeros are just padding so they're kind of wasted memory in in a way right um and so yeah that's that that's a big problem but with with um uh yeah with with the sort of the apis that that integated from flash attention you can just have a a long sequence and just have past the offsets right so you basically eliminate the padding and that means you can fit a lot more in the batch and and you know uh helps with utilization and things so um pytorch has a concept of Ed tensors um to be able to use those for this they have to kind of work they have to play well with some of the other things like compile um and uh and the pytorch sdpa scale. pro attention which I think they just so we you know this is um others in IBM research are sort of working on different parts of this like trying different angles of optimization so there's a there's some good talks from the conference um on uh in particular what we're doing around and and from uh from meta folks as well on pytorch compile um uh so yeah just I that was something i' forgotten to mention but um the pytorch version of this is called nested tensor so you know we may switch to using that now it's now they've just merged the support to make it work with these other things um this is just a picture of the uh the architecture of like our our version of this hugging face TGI that what we call tjs text generation inference server um it's it's actually uh a combination of rust and and python so all the pytorch stuff is in the the python side and uh you have a a the weights are sharded and you have a Shard a different process per per Shard um uh that's for the you know across different gpus so there's one of those per GPU uh and then the front end um is actually in Rust so uh rust is a super fun Lang I hadn't learned it before this and it's it's super fun language uh so um it's actually really good to to use for all the sort of batching and front end logic so that that's very fast and then every iteration it sort of updates the python uh processes over sort of internal grpc within within the same container uh sort of saying okay this is the next batch to and then the the tokens are returned back so it's you know the cycle is actually going back and forth between every uh and um yeah so that's overview okay so then fast forward a little a little while actually don't remember but it was a lot more recently uh a new project has sort of come up you guys may have heard of called VM this is from um some folks uh in Berkeley uh working with um like in the the uh Sky laab I think there um with working with yon stokker and they came out with a paper where they'd sort of taken a lot of these same ideas so the continuous batching um and they look they'd sort of gone even further with with looking at how to optimize the KV cache so a lot of how you optimize inferencing comes down to you know uh managing this cache because that's sort of often the the limiting limiting factor so and they had a really neat uh idea to apply sort of virtual memory techniques um P you know page tables and things like that to the the KV cache itself so the instead of like if you treat it as like one long sequence they kind of uh chop it up into blocks and they have some block size uh and then they have a block table and basically they allocate um they they allocate from this block table as the sequence grows so one of the problems uh if you don't do this is you have to copy the tensors for the KB cache on every iteration and that gets quite expensive I mean you either have to do that or you have to pre-allocate um but if you pre-allocate you're sort of wasting some space because that's okay I'm going to grow the sequence into this space uh and and often you don't know how long it's sort of indeterminate how long the sequence is going to run for right so it just depends um if there's a stopping an EOS token right so if the model decides okay it's just going to spit out one sentence or it's going spit out so you know a maximum but you don't know the actual uh so this kind of helps solve a lot of those problems and you can basically cram maximize the memory and make sure it's it's full with like real attention data and you're not having to reserve um there's not overhead of sort of fragmentation and and space that you have to reserve um but it was actually very comp complex for them to implement this because the um uh the kernels like the operations that do the um the work on the attention they expect sort of contiguous tensors um and obviously if they have a bunch of blocks all over the place uh they need to sort of deal with that and so they had to rewrite a lot of the kernels like specifically for page detention um to work with you know to understand this this block layout and then perform the operations on that uh but this project now has um has actually been gaining a lot of momentum they sort of committed to Apache 2 license uh so we're like now looking very seriously at this and um they have yeah lot very active Community lot of contributors so um I'd uh and they they're very um they're sort of associated with any scale so any scale like big big uh uh donators to them uh but they they're used with by other companies already um so I'm I'm getting close on the time uh I wasn't sure how long this would be but um I'll uh so there's a more huh okay yeah I I'll just quickly run through so there's some there's a lot of innovation in this space now so um when I when I was describing we were starting out and there was like these early things that this was like only yeah like I said last summer like 2022 um there was very few people looking at this but now there's uh there's like every day there's new papers coming out with different techniques to how you can maybe optimize and squeeze out you know different uh amounts of performance um very interesting one is speculate decoding this is where you have like a uh a much smaller model that is like a um uh I can't remember what it uh the the name they give it but it's it's it's much smaller so it runs much faster faster and what you can do is get that to predict a bunch of tokens ahead say like do maybe five predict five tokens in a sequence um and that because of the latency that is much faster than like one iteration of the large model the large model can then essentially review and verify those so you can you can pass in more than you know a sequence of tokens and and see like okay what would it have chosen uh and it can do that up to because all theit tokens are conditioned on the prior ones right so you know um you can get all the probability season one go if they were all correct then you know that the sequence the the large model would have predicted that and you can then continue if not then okay you have to like only you have to stop where they they got the next token wrong but then you can kind of continue from there so as long as it's kind of e like a lot of the text that comes out is sort of obvious like it's a the or you know whatever um uh so in those cases this this works very well um these are some of the other projects so there's um tensor R tlm from Nvidia is coming out tomorrow I think uh they it's like out in beta but they they've Incorporated a lot of these same techniques like page detention continuous batching I think they call it inflight batching for some reason they chose a different name um there's a project called flexflow from uh that's how they out of Stanford I think they have some kind of competition thing going on with the Berkeley guys that do the VM um yeah and then uh gml is like a more for like a is like one guy on his own that did did some crazy optimiz ations um that are very popular uh and here's a few other techniques uh won't go into them but um yeah if you're interested it's it's a very interesting space to get involved with uh there's a Blog that came out a few days ago from uh data bricks the Mosaic ml guys that also covers a lot of these things and you know and more and it's very well written and yeah gives a it's sort of a sort of scientific treatment to some of this um so yeah that's about it happy to take questions have a question to you I was curious how uh to hear more about the longer context windows and how that might affect I that's mostly during the training time and not inference in terms of the limitation but as far as serving people want longer context Windows how does that change uh is it is it what's the what's the limiting factor um yeah that's a good question um so yeah like you said models are trained um with a particular context in in mind and uh often if you go beyond that there sort of it degra even if you can have the capacity at the INF side that the sort of the quality of the output degrades um so they there's and that's a lot of that's because of the the way the positional encodings work so you um there's a bunch of techniques people have come up with to like extend that at INF without having to retrain indefinitely I mean one one of the problems as well like I did mention it with the the KB cache right so that sort of limits the sequence because it it grows um some of the techniques like flash attention meant that that instead of scaling quadratically it scales linearly so that sort of made it longer context more feasible uh there's other there's a really interesting one came out uh called um attenion sync which means you can do a sort of sliding window efficiently like that wasn't really possible before because um you'd have to like Rec the attention depends on all of the tokens so you can't really chop it off at the front because that kind of messes up the calculation right so to have a sliding window you'd have to recompute the attention for everyone but they found that you can like in have a a sort of sync at the beginning that and um and then have a sliding window so say your attention's 4K uh you can like you can just keep the KB cache and just chop it just truncate it essentially at the front and and you still get like good output and just you can just do that indefinitely then so you could have you know for input the context would still maybe limited but then you could have unlimited output um I think it's called streaming llm I think the paper great talk thank you just a a simple question so models like or the gp4 and Claud are they really a a team of models like there's like a look ahead model you're were talking about the the little model that looks ahead I'm just kind of wondering what if you have any if it's the same architecture as you're using or something like that I mean it's it's the same architecture up to a point but I think no one knows uh for some of I think there's speculation that that gb4 might be like yeah multiple models uh they there's this thing called mixure of experts uh which is a yeah where you can kind of combine M multiple smaller models I think that that get but yeah I don't I don't know oh you briefly mentioned about uh GG ml um how much did you look into that area um I haven't looked at it too much actually uh because yeah just cuz I I'd like to look more but I I had the impression it was more for people like trying to run um uh like the large models on their laptops right like on kind of consumer Hardware um and it like did amazing things for that so I I don't think it's necessarily targeting the sort of running on h100s and a100 gpus uh for like multi-user like all these concerns of like how you pack you know High very hike in currency um I don't think it's it's has that kind of focus so there are a lot of yeah ones that I didn't mentioned that are looking at optimizing for the sort of single batch size case for embedded applications all those kind of things as well um but yeah I I intend to look at it more so I think Oh there's uh all right there was one more question I can I off yeah