LLM Inference Arithmetics: the Theory behind Model Serving

Speaker 1 [00:08]

Hello, everybody. Super exciting to be here. Such a crowd. It's always a pleasure. Actually, my first time at a PyData conference in Germany, so this is exciting. Also, the microphone is really sensitive, so I'll try to make this a bit more ASMR. Okay. My name is Luca. In my day-to-day life, I work as an AI engineer, as we call ourselves nowadays, in extreme and I've been doing lots of stuff regarding to AI products of course and building a valve system but for the majority of the past year I've been working on helping our clients deploy models and use their own workshop like local resources to run inference on their cluster or kubernetes clusters and this was extremely enticing for on so many levels any left me on a huge like rabbit hole of trying to understand how a GPU works, spoiler, I don't know yet how that works properly, but that's an aside, and also the maps behind LLM inference. And this has brought me, as Vicky Boykes would say, closer to the metal. So today, we're going to talk about a bit more, of course, some duly disclaimers and go straight away to the takeaways and the key points I would like to bring you to and the real goal of this talk. We're going to see what kind of quantities we'll be interested in when we try to guesstimate how much, basically, VRM we need to deploy a model and which quantities we need to understand and their differences between training and inference. And also, we go through a bit of math if the projector allows. Oh, come on, there's not going to be so much math. Okay. Okay. Okay. I know that's scary. part here things we get a bit more complicated but if you bear with me they're gonna be fine finally we go to the juicy bit so the lesson there any conclusion you can draw from trying to deploying models on your own and self hosting them so let's start off with this since this is a lightning thought not another lightning talk but it's a 25 minute talk there will be some simplifications since we're already trying to present something that is It's based off this incredible talk about Sasha Rush. Sasha is also a professor at Cornell, is also now doing a reinforcement learning at Cursor. So he knows a couple of things about it. And he built this, realized this beautiful video called the Street Fighting Transformers. I highly invite you to check that out, especially if I miss some steps. The talk is partly modeled after his talk, and I do also recommend to check him out because the video is really beautiful. Also, I invite you to check out and mention the fact that this video he made is based off another book called Street Fighting Mathematics to learn how to do estimation problems with mathematics. Is it okay with the projector still? I don't think this is so much scary yet. But it's going to be really soon. So let's just state our goal clearly ahead. Model training and model inference are two different things when it comes to hardware. when it comes to compute and memory. In the end, these whole things boils down to one key factor. Transformers inference is one thing because inference and computing transformers can be done serially or can only be done serially doing inference, but can be done in parallel during training. These are some serious implications for how you deploy a model and all the kind of consideration that you can make. Maybe not you, you, but people that build the engines that allow you to deploy the models. So the mathematics behind this is really the key to understand. And we can basically prove this conclusion just going through the simple mathematics that goes behind trying to understand how a model behaves and the kind of requirements it has. The first corollary and the consequence of this is that we need to accept tradeoffs between latency and optimal hardware usage. How we do this and how we choose that, well, that's where the engineering part comes. And we need to build a sense of what kind of requirements that our product needs to have. But in general, self-hosting is viable when you can just do a sync batch processing, a prompt heavy tasks, which means prompts where you have long prompts and really short answers. So where generation is not really long. And we're going to see all of this why. So this is where we're going to end up with. Now we have a bit of math. let's get started we're going to talk about two dimensions training and inference and three quantities the number of parameters the memory in gigabytes and we're going to be a bit vague about how we mean memory because we can represent the numbers as floating points or integer floating points can be 16 bits can be 32 so we're gonna pretend that we are unitless for the moment and And the same goes for flops, because, of course, we basically, when we do learning, like deep learning, and we do matrix multiplication, but also addition. So, we're going to pretend that we do not care about this one. And, of course, the two dimensions that we take and we care about are training and inference. There's also fine-tuning, but we're going to pretend that doesn't exist for now. You can watch Sasha Rush's video about that. And he's also going to explain this for Laura, but we don't do that here. So, these are what we're going to analyze through the lens of a simple neural network. So, back to the drawing board. This is a neural network. Hello. This neural network has two parameters in and three out. What are the number of parameters during training? Well, it's fairly easy. We just connect the dots and we say it's Flipper-free. Yes, this network has six parameters. many parameters does it have during inference? 6i. Okay. He's paying attention. Congratulations, man. Okay. It has the same number of parameters. Okay. So, this is just to get started and get us warmed up. We can redraw this neural network in a projection notation, and we replace the 3 and 2 by D1 and D2, and so the number of parameters is just D1 times D2. Wow. This is super slow. Okay. and the same goes for inference this is fairly easy this is the easiest part of the lm inference arithmetics when we go to compute things start to change because this is our neural network during training well how many operations does it need to do to perform a forward step during training well of course it's d1 times d2 but also the batch size because remember that neural networks can be used to process and compute the multiplication in batch so you can operate in parallel on multiple inputs so the batch dimension is the first new parameter that we need to introduce in our explanation when we need to account for the arithmetic intensity of the tasks that we have at hand we also need to take into account the batch size and this is the focus step so if we think about inference well of course this is going to be the same thing right when we need to perform inference with the model yes of course we need to do the very basic same thing but training is different for one reason it's called training because we do the learning how do we do the learning the backprop that's it so we also need to take care about the back provocation step the computer requirements for backprop is approximately two times as much as the one that we have when we do the forward step why because we need to to propagate and compute the differences from basic computer derivative of the loss with respect to the input, which can be the intermediate input, and the derivative with the loss with respect to the weights. So effectively, and we start seeing it here immediately that when we do inference, the computer requirements are much lower in a way, and the training requirements are much more expensive in terms of raw numbers of operations that the GPU needs to perform. And this is basically this. So, if you want to be a bit more explicit about that, you need to do compute during training for the forward pass, but also for the backward pass. And I was saying before, we need to compute the derivatives of the weights with respect to the loss with respect to the weights, and also with the derivative of the intermediate inputs with respect to the loss, because we need to do chain rule. This is the second bit. Third bit, memory. At inference time, well, this looks kind of easy, right? we need to have all of the parameters in memory plus at least the batch size times the largest activations, where the largest activations means fundamentally just the values, the intermediate values that need to be materialized after you do the multiplication. So as a recap, number of parameters, fairly easy, just do the multiplication. Computing, if you're doing training, you do the forward pass and the backward pass. This means that it's likely three times as much compute operation compared to the case when you do just inference. Memory, you just always need to have the number of parameters in mind. When you do inference, though, you need to also take into account that you need to materialize the intermediate results of the multiplication. You need at least to have enough memory to store the biggest layer in terms of activations because you can just delete the previous steps. We're going to pretend the skip connection don't exist. Well, approximately, this is street fighting transformers, so we don't do things perfectly. We're just trying to get a sense to guesstimate how much memory is required. For training, well, story complicates a bit, right? We always have to have in memory all of the parameters, of course, but we need to do backprop. So we cannot just store the largest activations. We need to store all of the activations. And if we want to win an ICLR test of time award for a paper, we will be using the Adam optimizer. And we need to preserve the state of the optimizer in memory. And usually, I'm not going to the details of the Adam. You can read the paper, which just won an incredibly important award. It's basically just two times as much parameters as a number. Two times as much as many parameters as the total parameter count of the model itself. Because of how the optimizer works. So we see that there is a dramatic shift in terms of memory requirements between inference and training in this really simple case. Because at inference time, we only need to care about basically the biggest layer. Whereas during training, not only do we need to have enough room to have all of the intermediate activations so that we can just do back, but we also need to maintain the optimizer state. And this is a lot. So, we just saw what the difference is. And you'll be like, okay, Luca, this was the simplest thing you can come in with. But we have to talk about transformers. We're not doing bait and switch, right? Well, now for the toughest part. If you think about transformers, this is the regular transformer. This is a picture from the original paper. The original paper is encoder, decoder. We don't really care about this. A transformer is mostly like chat GPT, decoder only architecture that's made up of three building blocks. The input embedding that has a size of V. And then N layers that are basically made up of a feed forward layer and the attention mechanism. Let's just pretend normal layers do not exist. So let's talk about them in terms of memory, number of parameters, and compute. The embedding layer is the easiest one. Because it is kind of big, V is a vocabulary that's basically a lookup table that maps IDs, integers, into input vectors, basically vectors. And this can be quite big in terms of number of parameters. Usually two to three orders of magnitude more than the hidden dimension of these models. So it's V times D, where D is the dimension that it maps to. Basically like, I don't know, GPT-4-0 should be 200,000 in input. The compute, though, it's negligible, because D is just a lookup table. Memory, well, you just need to store the vectors in mind in your memory, plus the optimizer state if you're doing training. So, this is fairly easy. The feedforward layer, we just saw it. In the original Transformers paper, we have this kind of structure, which is basically a double MLP, MLP into the other one. And we know how to compute that. Number of parameters is N feedforward layers times two times D times 4D, because all of these layers are connected. Compute and memory are the same ones as following the rules that we outlined below before. Now for the attention layer, because this is a bit of a beast. The attention layer is made up of a first layer that separates. I'm not going into how attention works, but I'm going to assume, right? And I'm doing math really well for this now. That basically splits the input into three, query key and value, the attention mechanism that I'm going to just represent as a squiggly line, and then finally an output projection. And you see here that if you look closely, if you squint a bit, our transformer is actually mostly made up of MLPs. Because outside of the embedding layer, but whatever, and the attention mechanism, insofar, the whole structure of the model was just MLPs, right? So we know how to compute a number of, we know how to estimate the compute that it takes, we know how to estimate how much memory it takes, both during training and inference. So let's go a bit deeper into the differences. Well, if we want just to compute the number of parameters that we have in the attention block, well, it's just N times D times 3D plus D times D. That's easy. Memory is just the same as a simple neural network. How about self-attention then? Well, this is where it gets interesting. Because attention doesn't have any parameter. But it has a lot of compute. The compute is B times D times D. Where T is the length of the sequence. Because attention looks backwards. So, you need to compute it over the whole content that comes in. It can still be parallel. Because if you accept multiple requests, so, multiple sentences, multiple sequences, then you can compute attention over all of the sequences as a whole in parallel. Provided you have enough RAM. So, no parameters. compute requirements remember the t is going to be important as far as compute goes well we kind of know how to deal with this as a recap right of the whole attention block we need we just have d times 3d the forward step during training plus t times d because we need to compute the attention on the whole structure on the on the whole sequence plus d times e where it's the out projection simple mlp plus the mlp like the final mlp lady feed forward layer And we need to multiply this times the number of layers that we have times the batch size times the sequence length because we need to generate one step at a time to generate the new token, right? And if you squint, you will notice that these T are actually coming out if you factor out the equation at the T squared. And keep this in mind because this will be important. If we're doing training, we just multiply this by three because we need to do the forward the backward pass but ideally this is just the same between training and inference for when it comes to memory we just need to store the number of parameters times the number the layers of the biggest activation and the optimizer state and the old activation sorry which i'm writing down here plus the state for the optimizers that's okay when it comes to inference it's just what we saw above so d times 3d because we need to do the forward pass plus t times d which is the attention mechanism, plus D times D, which is just this out of projection. Then we need to add the inference of the line as fifth or layer. We need to multiply this by N times D times B. But this is where things get complicated. And this is where it comes that attention is quadratic. What does it mean that attention is quadratic in compute? It means that if your sequence length goes from 100 to 1,000, it's not like you need 10X more compute. You need to have 100X more compute. Because Because a long-wear sequence needs to attend to every previous token on behind. And this is where the behavior, the quadratic behavior pops out. In theory, this is the same between training and inference. But during training, you can actually make this more parallel with a simple fact. When you perform training, you can just generate one token at a time. And you know the whole data set beforehand. You know the whole sequence beforehand. You're trying to learn how to predict the next token. So basically, you can be clever about how you compute the attention, and then you can just spread out all the training codes so that every example can run in parallel efficiently. When you do generation at inference time, when you ask ChatGPT to build, like, whether you should put glue on your pizza, it's generating new tokens one at a time. So it cannot be parallel. It needs to be serial in that kind of generation. And this is the source of quadratic behavior that complicates a bit of stuff. And this is where the novelty comes out. Because since attention looks back and requires Q of the current step times the K and V, the key and value with the size 2D from the past T minus 1 steps, instead of recomputing all of this, we can actually store all of that into memory somehow. And this is actually really clever because if you think about it, attention is squared in terms of complexity in compute, but if we needed to be doing that for every single new token, it's let's just pretend this is serious, but it's just cubic because you need to compute the attention at step two for step one, at step three for two and one, at step four for three, two and one. And any time you just do one step more, the computational complexity scales. So you can just say whatever, K and V are static, let's just compute all of them beforehand and just store them in a cache. That's the KV cache for you. And so basically we are trading compute for memory. We are making our memory requirements higher so that we can have faster compute. I said faster, though, not more efficient. But let's just not take this step too far. So memory requirements for transformer inference, you take B times the largest activation, the the number of parameters, of course, plus the KV cache. The KV cache is the first two, basically 2D, so basically the first step of the attention mechanism, times N, the number of layers, times B, which is the number of batch, of samples we have in our batch, and times T, which is the length of the sequence. And this KV cache is the amount of memory that you usually keep being seen around as a 20% of the model size. So basically, if you have a Mistral 7B, you will likely have a seven times two, which is the memory in terms of number of parameters. We tend to ignore the largest activation. And then we add that 20% of that size, which is usually the smallest amount that you need for a very small KV cache. And this is where the number comes out of. And one of the discoveries that I made last year. Inference engines do implement a really more serious approach of these KV cache, which is fairly simplistic. But in a way, the end goal of all of this is, as I was saying at the very beginning, there are clear differences between how the attention mechanism can be parallelized during training and inference. During training, it can be computed more efficiently. During inference, we need to generate new tokens. It needs to be, it can't but be serial. So you need to reduce the amount of latency and thus compute. We can just store the intermediate values in a KV cache. This, however, is still suboptimal usage of the hardware because you need to move in and out the data from the memory. You need to move them from the part of the GPU that has the memory into the part that has, that does the computations, the ALUs, so the Arithmetic Logical Units. So this movement is actually building up an IO bottleneck that slows your generation down anyway. So while we can be a bit more efficient in reducing the latency, we're still underutilizing the hardware with respect to the training. So how do we do this? Well, wow, this wasn't here. Let's just take a step behind. When we factor in KV cache, generation actually becomes split in two steps. And this is really fascinating. Because in the first moment, we need to fill the whole KV cache. This phase is called pre-fill. This step is a matrix-to-matrix multiplication, which can be dramatically fast. And this is actually the point in the generation where the usage of the GPU of the hardware itself is maximum, is optimal. However, when we go to the decoding phase, generating token one at a time, we encounter the inefficient that we were talking about. Because even if we have the KV cache, we need to read in and out the data from the KV cache. Basically from the memory of the GPU back to the computational units back and forth. And this bottleneck slows down the GPU, which can sit idle in while it waits to perform a new multiplication. How do we deal with this? Well, people have found out that they can just say, okay, increase the batch size. Because if you have four, five, six more requests at once, you can read at once all the values from the KV cache that are required for that step, perform the generation, and drop them when you're done. But this has some problems. When it increases the efficiency, the arithmetic intensity of your GPU, it increases the number of tokens per second that you can generate, it decreases the time to first token, which means basically when you get the first bit of a reply, which is by far the least interesting most of the time. It's like, yes, of course I can help you with that. And this is decreased, it increases the time to first token, sorry, this was a typo. It increases the time to first token, that's my bad, because we need to compute the KV cache for more tokens at once, for more sequences at once, right? And these also have some other limits. Because we cannot just say, okay, I'll just use a batch size of five. Because if you cannot start a job until you have five requests, people will be hanging. So you need to implement a more clever implementation of batching. And this is what Triton by NVIDIA, BenchML, and many other more inference servers implement. But still, even if you have a clever implementation of batching, the throughput doesn't increase linearly. So, even if you make the batch size optimal, like bigger and bigger, you cannot expect linear gains. At some point, you will start hitting diminishing returns. Because of what? Well, of course, because of memory bandwidth, you still have to move data in and out from the VRAM. And if you have more batches and more samples in your batch, you need to move in more and out data out from the KV cache. And also computation limits because it might just incur an out-of-memory error. So all of this is when you deal and anybody deals with when they're trying to self-deploy their models. Should you ever want to do that, well, as we saw from the structure, if you do batch jobs where you can control and perform them asynchronously with an optimal batch size, you might find yourself that this is a perfectly fine use case where you can save money. If your completions are really short and you have a huge pre-filled phase, but just a short number of tokens that needs to be generated, that's where self-hosting is a viable solution. For any other task, run your own experiments. In a way, OpenAI has such a massive scale where they might just be getting thousands of requests per second, so they have, like, groups of eight, each 100, and where they can optimally deal with all of this, stuff that we cannot deal with. And so for those occasions might be just best for you to just rely on the old curl OpenAI or any other model provider. And this is it. I will thank you very much for it. I would love for you to scan your QR code and share any feedback about this talk. This is the first time I give this. And I'm really excited about it because it's a topic I've been working on for a while. I would love to connect and get in touch about how this works, how I can improve this, and I can help you out with this. I know this was a bit rushed, so if you have any questions, feel free to drop by and message for the rest of the conference. If you need some more reference, I have them here. I will share the slides somehow, so you can just click on the links. But do check out Street Fighting Transformers, all of the others, our blog posts from NVIDIA, BenchML, and also the guys from Cursor. The last one is really interesting to read. Thank you very much.

Speaker 2 [25:26]

Thank you very much. It was a very energetic talk. I think you somehow lost me and maybe others too. We have one question on Slido at the moment.

Speaker 1 [25:37]

the moment so one more than i expected

Speaker 2 [25:39]

Feel free to add more while I read this and Luca is answering it. So could you say something about different inference engines, OLAMA, VLLM, TENSOR-RT and when to use which?

Speaker 1 [25:54]

So, Ollama is meant to be used as a wrapper around Lama CPP, which is a really popular framework that is designed to run servers on your machine, basically on CPU. So it's designed for you, I can't say mostly for local use, but in constrained hardware. When you use Ollama, actually anything built on top of Lama CPP, which means Ollama LmStudio, you will likely need to be working locally with quantized models. VLM, SGLang, and TensorRT are two different things. Sorry, they're three. TensorRT should be, in theory, one of the most performant ones, but it's really hard to get right. So, it really requires a lot of tweaking and fine-tuning, understanding the documentation that's, I can't say poorly written, but it's not just as elegant as other, as intuitive as other frameworks. But, on the other hand, I think you are going to be so, like, unless you go out on a real high scale, you don't really need to focus too much on which one is the best. Just pick the one that you can deploy more easily and then start running measurements. And then you can start investing more time in finding out which is the better one. So in a way, like, TensorRT is both a compiler, so basically that optimizes the structure of DLLMs and has really made kernels that allow you for your models to run faster. But if you need to configure it yourself, it's going to be really hard. It's really designed to be optimized with NVIDIA hardware. So it should squeeze the biggest amount of performance. If you work with VLM or SG-Line, they're in Python. Basically VLM is just it has custom kernels, but it's also basically a fast API back end, so with a lot of incredible stuff on top, of course. But nevertheless, you can deploy more easily. You can also just run VLM serve something, and it can be your, I don't know, Docker file entry point, and you can go with that.

Speaker 2 [27:47]

So, perfect. While the audience is thinking of more questions, I have one as well. Please do. So, I understood that if you do batch inference, the time to first token is increasing because you need to move more data from memory to compute.

Speaker 1 [28:03]

Actually, I wasn't clear enough with that. You need to compute the KV cache for all the sequences. So if you have ten times more sequences, it should take a bit more time than computing for eight. So this is why you decrease. Because you need to end the profile step for every sequence before you start generating your stuff. So if you have, like, three times as many samples in your batch size, of course, you're going to be waiting a bit more time. Also, especially if the sequences are heterogeneous. they might be just some of might be longer than the other ones

Speaker 2 [28:32]

Okay, and once you moved all the data from memory to your GPU Then it's the same time to compute all the batches in parallel

Speaker 1 [28:41]

in parallel? Yes.

Speaker 2 [28:42]

Okay, how is that that's some kind of magic for me that like it's possible to compute more

Speaker 1 [28:48]

So, basically, every step takes just the number. It should be basically equal across all of the sequences. Of course, some sequences may be longer, so you might be just asking to summarize a piece of an article, or you'll just be doing deep research or something. So you need to be clever how to balance the jobs so that jobs with approximately similar length run on the same hardware. If you have Google has GPUs that need to run deep research, those GPUs need to be on their separate class and it cannot be just mixed with requests that just are designed to just be answering with a yes or no question i'm doing extreme simplification but the thing is you need also be able to balance the kind of requests that hit a certain machine compared to others

Speaker 2 [29:30]

Perfect. Thank you very much. Thank you to the audience. And now you have 10 minutes to get to the next talk.

Speaker 1 [29:37]

Thank you.