Reiner Pope – The math behind how LLMs are trained and served

So maybe this will help motivate the kind of analysis that you'll be doing through the live share. Great. I mean, to jump to a little bit to jump to the conclusion, the big effect is batch size, but what what we're gonna do now is quantify exactly what that looks like and what its implications are on latency and cost.

There's gonna be another effect, which is you can call it speculative decoding or multi token prediction. Can We maybe come back to that later, but I think the first thing that we'll talk through is batch size. So what I'd like to introduce is sort of the two principles of analysis.

Firstly, we're gonna look at a roofline analysis of how we run a transformer model on a cluster of chips. We'll take a sort of a, let's say, a Blackwell NVL 72 cluster, so a rack of 72 GPUs. And so the roofline analysis means we look at memory bandwidths and and compute performance.

And then the other side of that is that we're going to look at just two simple factors of the model, which are the time to operate on the weights and then the time to operate on the context, the KBCache. So let's jump in. What we're gonna try and do is we're gonna try and estimate the time that it takes to to run an inference of a certain shape.

Now we're not perfect here. We can't exactly predict the time, and so instead we're gonna approximate. And so we're gonna say that the time must be greater than or equal to a certain quantity.

And so we're gonna consider two different aspects. We're gonna look at the time for it takes to do the memory fetches, and then the time it takes to do the compute. And it'll turn out that this actually gives us a very strong predictive power even with a simple lot.

So one by one, what is the time that it takes to do the compute? So there are really two things I need to do in the compute. I need to multiply by all of the active parameters, and then I need to do some work on the attention.

So multiplying by all the active parameters. I have a certain batch size that I'm running, and then I've got a number of active parameters in my model. And then and then I'm just gonna divide this by the compute throughput, which is the FLOPS of the chip.

So this is a hardware concern. So this this actually accounts for all of the compute time for all of the weight matrix multiplies. There's a little caveat here.

We we've sort of ignored the time to do any of the attention computation, but that in general can be will be quite small in comparison to this. So so we'll ignore this. Maybe I'll just interrupt from time to time to ask some very naive questions or to clarify some basic points.

I mean, we will see exactly why batch is such a favorable optimization. But what will turn out to be the case is that if you do not batch together many users, the costs and the economics you get can be like a thousand times worse than if you do batch many two users together. And we'll be able to see that quite especially.

And then number of active parameters. This is saying, if I look at, for example, a DeepSeq model, the DeepSeq v three model has about 37,000,000,000 active parameters and then 700,000,000,000 total parameters. So this is we're we're focusing on just the ones that are active for a single token.

Okay. So where model compute performs, I'm gonna keep writing equals, but in in all of these cases, can think of this time as being at least this much, and and maybe there'll be some terms we ignored. On the memory side, what do we need to do with memory?

We we need to fetch we need to fetch all of the weights. And so there is some time to fetch all of the the total number of parameters, not just the active parameters. So there's wait fetch time.

And then in addition, there's a kvcache fetch time. So there is this actually depends on batch size. So for every element of the batch, we have to fetch an entire context length worth of tokens.

And then there's a size per token, so, like, bytes bytes per for for one token.

what the KV cache is real quick. Yeah. So when I do a forward pass let me draw actually a how the autoregressive inference works.

So this is during decode. So if I think I have a bunch of tokens of text, I'm growing a tensor because, ultimately, the tokens are represented as some, like, tensor of in some embedding dimension. And then in this direction, I have the sequence length.

The work of running a decode is I I have to run each token through a through a whole bunch of matrix multiplies over a bunch of different layers. And I have I have in general, I'm gonna have to do that work over all of these tokens. But then one step of decode is actually to produce just this one additional token up here.

Yep. And so what I'm gonna do there is I'm gonna run a full forwards pass of multiplying by all of the weight matrices in the entire model. But then I've got this attention mechanism where this token sort of it's it's like looking at all of the past tokens in this way.

And what is it looking at specifically? It is looking at some internal representation that the model has produced of the tokens, and we call that the KBCache. So this process of attending this this single token attending to all of the history of tokens, that's attention.

It is mostly dominated by memory fetches rather than than matrix multiplies. Mhmm. So we've got the amount of memory that we're fetching, shown over here.

And then there's, of course, just then divided by the memory bandwidth. So so the memory bytes per second. So, in fact, these equations here are actually enough for us to now draw some fit lines.

And so the things that we'd like to look at are sensitivity to batch, and then also, which we'll draw separately, to context links. So we said that the big big effects you can get is, like, some some trade off in latency versus versus cost in in batch size. So so let's draw them out.

I think there's just really two graphs we want to draw. We'll first just draw batch size versus time here. So when we look at the shape of this, we've got a maximum of, well, the sum and then and then another term.

So let's look at these terms one by one and how they scale the time for compute and and memory and how they show up. So let's first look at this compute time. This is just purely linear linear in batch size with no no offset.

So it is some kind of like this. This is this is t compute. And then on the memory side, we've got some portion here that that is just this constant that's that is, you know, constant in some base offset here, which is the wait fetch.

Wait fetch. And then finally, have this term here, which is the KB fetch, which we're gonna draw as as there's the KP fetch, which is which are linear in batch size. So it looks like that.

So the sum of this plus this maxed with this. So let's at least first draw the sum. So the two memory times in in conjunction end up looking on this curved slope like this.

Mhmm. And then we get a the overall maximum is I'll draw a little thicker here. Is it a maximum of these two curves?

Makes sense. Okay. So so so what what does this mean, actually?

So this is a latency plot. So if I grow my batch size, I get initially some not very strong dependence on batch size. And so there's some lower bound on latency here, latency lower bound.

Lower bound. So this already partially answers the question. For a given hardware configuration and then we can talk about varying hardware configuration.

But for a given hardware configuration, there is a lower bound on latency, which is simply the I need to read all of my total parameters from from memory into the into the chips. And that takes a certain amount of time. If if I use all of my memory bandwidth, I can't do any better than that.

And what implication the KB has on memory

So I think we should come back and explore this. Yeah.

the compute time? Yeah. Whenever we have balance points, it kind of says that you're getting it exactly right.

And so for the particular context length where the slopes match, that says I am equally memory bound and compute bound Right. Which is a really desirable place to Yeah.

There's a key point here that I'm modeling this context length as or I'm modeling the memory fetch as linear in context length. That actually depends on model architecture. It is true for many of the all of the model architectures with dense attention.

Yeah.

I'm pretty excited about sparse attention. It's hard to know what the labs are using. DeepSync has published a sparse attention mechanism.

I'll just put a plug in that sparse attention some of the DeepSeek papers that have published sparse attention end up putting a square root in this term. Okay. So so far, we've looked at the latency.

It's kind of hard to read off cost from this. So if I think, what does cost mean? I'm gonna like, to run this inference, I'm gonna use the GPU for a certain number of seconds, like one millisecond or twenty milliseconds or something like that.

And I have to pay the rental time for for that that time. So, like, it's $2 an hour per GPU or something like that. So so that's the cost of this inference.

But how much value have how many tokens have I processed during that inference? That is the batch size. And so what we actually wanna plot is gonna be the the cost versus batch size, which is like t over b versus batch size.

This is the cost per token. So, like, we have to imagine dividing each of these three curves by by b, so multiplying by this reciprocal. And so what we end up with there is the the compute curve is gonna it was linear.

We divide by b. That makes it a constant here. And this is t compute.

The k v fetch was linear. Now it becomes a constant as well. K v fetch.

And then the the wait fetch was constant, and now we've divided by b, and so it becomes this hyperbolic. Mhmm. And so, again, we're gonna compute the max of the sum.

So the sum of these two terms shifts the the parabola up. The sum of k b fetch and the weight fetch gives us a sort of a higher parabola that's like this. Mhmm.

And then we're gonna take the max with the compute here. So we end up with this this being the overall shape that that we care about. So, again, so, like, we see some limiting behavior.

The cost initially starts very high at batch size of one. Actually, like, it almost goes to infinity. Like, it's, because we've got so many weight fetches which are not amortized over a large batch size.

But then as we increase the batch size, the weight fetches become amortized over so many different batch elements that they their cost goes grows very small, and eventually, the compute time, ends up driving the cost. Mhmm. So there is a limiting, like, lower bound lower bound on cost, which is this one here.

Yep.

a much bigger batch. Yeah. Yeah.

They're unique per batch. The compute is also unique per batch. And so what is the minimum work you can do per batch after amortizing everything else away?

like, yeah. How big are the batches Yeah. Practically for frontier models?

You can you can just solve for that, actually. And it's not even particularly sensitive to model architecture. So let's let's go ahead and do that.

So what we're talking about is we're gonna say when the memory time is equal to the compute time. That's that that's what that question is. For now, I'm gonna discard the because we're focused on what what the batch size is, and, really, there's a question of what when the weights are amortized over the Yeah.

The the multiplies, I'm gonna focus on comparing the weight fetch time to the weight multiply time. I'm gonna disregard the KB fetch term just just to simplify the analysis so we can get a kind of a clean answer out. So we're gonna equate this portion with this, with these two terms.

Yep. So writing that out, we get n, number of total parameters over memory memory bandwidth is equal to batch size times number of active parameters divided by the compute performance. So looking over here, everything on the top, these are model parameters.

Everything on the bottom, these are hardware parameters. It it turns out to be nice to rearrange them such that we have the hardware parameters on one side. So so that's this is equivalent to well, so the memory bandwidth being equal to batch size times number of active parameters divided by the number of total parameters.

So so this is a hardware parameter. Actually, the this actually ends up being a dimensionless constant. If you look in terms of FLOPs, what are the dimensions of this?

This is multiplies per second. This is bytes per second, so that's not quite dimensionless. But what you do is you say, like, multiplies per second times let's say I'm doing f p four.

So I I do, like, how many f p four multiplies per second times the fact that each one each f p four is half a byte. And so I can actually make this end ending up being dimensionless. And and this ends up being on most GPUs around 300.

To what extent has the hardware changed? So from, like, a 100 to a 200 to b 100, the the flops has increased substantially. The memory bandwidth has also increased substantially, and it has remained reasonably stable.

Yeah. And we can we can express this one as well. This is a sparsity parameter.

Yeah. And I I might even phrase it slightly different. Let's solve for batch size in total.

We end up with I'm sorry. Just moving this back over to the other side. We end up with batch size needs to be bigger than approximately 300 times sparsity.

So, for example, if I have a 100 like, I activate in DeepSeek, I activate 32 out of 256 experts, so this would be, like, eight for DeepSeek. Got it. Okay.

So so this actually gives you a ballpark, which is, like, remarkably accurate to practice. Generally, people will go a little bit larger than this. They don't really want to be exactly at the balance point because, real world efficiencies aren't as good as a a refinement analysis would say.

Should grow larger. So this is gonna like, we we solve for the equivalence between when compute time is equal to memory time. If I add in more memory bandwidth, like, something that consumes more memory bandwidth, then I have less available for the for the workloads.

And so I need to grow the the memory bandwidth more, the the and therefore the batch size more. This seems incredibly small. Like, a batch this would be, like, less than one sequence.

Right? Yeah. Okay.

So so I guess this is keep in mind that I'm talking about the number of tokens that I'm generating one more token for.

it it would not take you a hundred milliseconds to fill up the next 2,000 slots. Yeah. The the way to think about this, I guess we think think of it as like when does the train depart as a model?

Yeah. So let's say I've picked a batch size that I'm gonna run at. Maybe I pick this batch size.

And by the way, this intersection point is the same intersection point here. So picked this batch size. I know that it's gonna take, for example, maybe it's something like twenty milliseconds is a common place to sense up landing.

Mhmm. What I'm going to produce is like, so this is a timeline of what is running on the GPU. It's gonna start a new batch every twenty milliseconds regardless.

And so so each of this is 20. This is 40. You can think of this as a schedule for the train.

A new train departs every twenty milliseconds. Any passengers who are ready board the train. If the train is full, then they wait till the next train.

If the train is not full, the train's gonna go anyway. And so in terms of what that means for queuing latency, it means that the worst case is that you like, a request arrives just after the train departed. It has to wait for the next train, so that's up to twenty milliseconds, and then it has to wait for that train to to complete.

Yeah. And so the worst case latency is forty milliseconds. So how is the train milliseconds derived?

I mean, rule of thumb, but where it comes from is not fully explained yet, but so far we've focused on memory bandwidth and compute time. When we look at memory, the other consideration is that we want to use all of the memory capacity we have. And so, generally, we're going to use all of that memory capacity to store the weights or the kVs.

And so we just want to read, like, in the time of doing a forward pass, maybe we want to read all of the memory capacity into into the chip. And so that is capacity divided by bandwidth. That tends to be twenty milliseconds on on many gen different generations of HBM.

byte divided by bytes per second. Yeah. So, for example, I mean, on on, I think, the Rubin generation, it is something like 288 gigabytes, divided by 20 terabytes per second.

And this looks like it comes out to about fifteen milliseconds.

if I think what that means, it means I actually have time to read the HBM, like, twice. Yep. By the way, most of HBM accesses is reads, not writes.

It's like almost all reads because the weight matrices are read only, and then almost all of the KV cache accesses are So in like, let's say around thirty milliseconds, I can read all of HBM twice. But what's the point of that? Like, I I don't wanna read the weight matrices twice.

push towards centralization is it that you would have these economies of scale from inference, from batching? Yeah. But it seems like it's not that big a deal.

Like, I don't know. Two is 2,000 users at the same time a lot? It doesn't seem like a lot?

So we can do a bit of analysis on this, which would be actually it's like you can think of it in terms of number of users, but maybe a more productive way to think of it is in terms of number of tokens per second. Mhmm. So what does this batch size mean in terms of tokens per second of this of the system?

So tokens per second. Tokens per second is going to be equal to the batch size. We run a batch many tokens, and then we do that every t.

So So every time intervals, which is, let's say, which is which is this thing is equal to the fifteen milliseconds twenty milliseconds number. So this ends up being batch size itself times about 60. So, like, 64 times b.

And so this ends up being around 2,000 times 64. So, like, a hundred and twenty eight hundred and twenty eight k tokens per second. So this is sort of in more digestible units.

Like, it's hard to reason about concurrent users, but what is the track global traffic for for a system? When you look at some of the announcements, sometimes the API providers will will brag about how much traffic they have. The the the numbers that I've remembered from some announcements of Gemini last year were in the hundreds of millions of tokens per second worldwide.

Yeah.

your saving compute by increasing the sparsity factor? Yeah. So perform equality of the of the model, you you mean, rather than speed of the model.

Yeah. Yeah. So, unfortunately, we're not able to answer that analytically.

That is an empirical question of model quality. Best I can do is pull up a paper and answer that empirically. Yeah.

Okay. Should we follow the paper now, is that makes sense yeah. So this paper, this is Unified Laws for Routed Language Models.

It's a somewhat old paper by this stage, but one of the things that they did is looked at, if I keep increasing sparsity, what is the model quality impact? This answer is very sensitive to the actual choice of Mixture of Experts. Mixture of Experts has been around for a really long time.

I think it was maybe even back in 2017. But the tech techniques have changed a lot. DeepSeq mixture of experts was was a big change in how it worked.

There have been older papers which are G Shard switch transformer. So the actual empirical results are gonna depend on all of that. But on one of the older techniques that is shown here, you can see if I hold constant the number of active parameters at a certain size, and then I increase the sparsity, which they call expert count here, The quality keeps increasing.

And then if you imagine drawing a horizontal line from 1.

parameters. Yeah. Mean, actually, even more so.

Yeah. A huge increase in parameter count for a modest increase in Yeah. So in this case, it's what what is it?

Four x? 64 x for four x. Yeah.

is that good or bad, actually? Even from a memory point of view, keep in mind, you are doubling this portion of the memory fetches, which is amortized by batch. And so just just keep running a larger batch size.

From the point of view of the the analysis we've done here, this is pure win. Keep doing it. Keep doing it until you run out of available users, basically.

Mhmm. So there's actually this equivalence between if I want to go sparse, or if I have a lot of users, I can go to a much sparser model. So from that point of view, it's it's a reasonable trade off.

Yeah. So To have more sparsity. Yeah.

So, I mean, maybe this would be a good point to actually talk about how a Mixture of Experts layer is typically laid out on on a, like, on a rack of GPUs or something. Yeah. Yeah.

Yeah. Makes sense.

Sparse Mixture of Experts. Yes. Maybe how we lay that on out on a GPU.

Yep. So let's zoom in on the Mixture of Experts layer first and and and sort of draw what that looks like. So we typically will have a some kind of a router layer Mhmm.

Which is making the decision of where we route the experts the the tokens to. So we get tokens coming in here. They go through a router layer, and then we have a bunch of different experts.

I'll draw draw a few more to line some up. And then the router will make a decision on which experts am I gonna route to, and it'll be a small fraction of them, maybe one in 32. So maybe it'll make a decision to route to this one, maybe this one, and maybe this one.

These experts so these each expert itself is a normal MLP. It has a up projection and then a down projection and a nonlinearity in between. And then finally, we sort of do the inverse operation.

So where we were broadcasting things out here, we're gonna bring them back in and sum them up. So bringing them in like this. And then finally, we have our residual connections that the token is also passed through here, and it gets added to the result of the MOE layer.

So so this is a normal MOE layer. What I want to talk through is how this is mapped to a, like, a GPU rack and what this means for communication. Because I think this will will start to show some of the the limits of how fast we can go.

Yeah. So the standard practice here, and it it is the best solution, is to use expert parallelism. So that means different experts go on different GPUs.

So if we take something like a DeepSeek model, they have 256 experts. Let's say we want to run that on a Blackwell rack. So there are 72 GPUs.

We have a divisibility problem. This is not a power of two. So we'll just, like, simplify and say we're only gonna use 64 of them.

Just ignore the other eight. It's not a big deal. And so we we have four experts per GP.

Very simple. For the sake of the diagram, I'll actually just say, let's let's say we have two experts per GPU. So we we end up just putting these are the GPU boundaries.

Every pair of experts is on its own GPU. And then we can look at the communication cost. We had some experts stored there's some tokens stored centrally here.

They get routed to all of these experts. And so there's some communication cost paid here. There's the same communication costs paid on the output.

And then the hope is that this does not become communication limited. Now, what is the traffic pattern here? The traffic pattern here is that any GPU, in fact, will be talking to any other GPU depending on the decisions made by the model.

So this is an all to all traffic pattern.

the router is more than one GPU? Yeah. The route so I I drew this as one router.

In reality, you would actually have many copies of the router. And so you would have as as many routers as as GPUs, in fact.

Yeah. So these are 64 GPUs. These are 64 GPUs.

It's actually the same GPUs. We just draw them as separate because they're serving different purposes. So at this point, any GPU can be sending to any any other GPU.

So this all to all pattern of communication that shows up and how how the Blackwell Racks are configured is a is a perfect fit for the the communication pattern that the MOE actually wants to do. However, if you think, maybe I want to do like, maybe one rack is too slow and I want to do two racks, then I have this challenge that, like, maybe I've got some sort of rack boundary drawn outside here like this. And I no longer, in fact, have all to all communication between all the GPUs in two racks.

And so the rack to rack communication ends up being a substantial bottleneck. So this sort of like, the fundamental thing here is that one rack is actually the bounds the size of an expert layer you can do. And so this has been part of what's been driving towards larger and larger interconnect domains.

Yeah.

and within a rack Yeah. And the all to all versus not all to all nature of communication within versus outside. Yeah.

And and this is a place where it starts to be very different, in fact, between NVIDIA, for example, and Google, and then others, including us. The so generally, a rack is a it is a physical structure. It it's a few meters tall, a meter or two wide, depends on configuration, and it stores some number of GPUs or XPUs, which is typically about 64.

The the cons what constraints it being a certain size is power delivery, weight, and cooling ability. It ends up being about this size in in many cases because of these physical constraints. So then when I deploy a data center, like, I've got a data center may have thousands of these racks.

So I've got one of these tall racks that's got a bunch of GPUs in it and so on.

In NVIDIA's case, the the communication topology is actually they put the GPUs on on the outside of the rack, and then they put these switches on the inside of the rack. So what this ends up being is that there's a set of switches in here. These are the NV switches.

Mhmm. And then they run a bunch of cables. Every single GPU has cables going going to the switches in the middle.

So every GPU goes to switches in the middle, and then the switches have connections to all the GPUs. So all of the GPUs can talk to all the other GPUs in in just, like, two hops, going to the switch, going to the other GPU. Now when I want to leave the rack, I end up going via a different path.

The GPUs have also a much slower connectivity, which is typically about eight times slower, which is so so the green that I drew here in GPU cases is the NVLink. More generally, it's called the scale up network. This is the scale up network.

You will typically also have a scale out network, which allows you to connect to, like, some data center switch. So data center switch. And then all of the GPUs will have some connectivity up to some data center switch somewhere.

Mhmm. But this is this is about times like, this is the scale out, and it tends to be about about eight times slower in in bad words. So the challenge, if you want to, for example, lay out a Mixture of Expert layer across two racks, is that half of the GPUs here are gonna be wanting talk to the GPUs here.

And so, like, half of the like, just on average, like, when I look at where the tokens on on these GPUs want to go, half of the tokens are gonna wanna go inside the rack. That's great. They can use the the fast scale up network.

But half the tokens are gonna want to leave the rack and go to the other rack, and that's not as good. They're gonna need to use a much slower network. And so that becomes the bottleneck on on the all to all pattern.

A different choice would be, well, why don't I, like, have a big switch here and sort of, like and connect everything to some big switching like, much bigger switch that actually combines the two racks together. There are many ideas in this direction, but in general, it becomes the reason you have this sort of hierarchy of switches rather than one big switch is to manage the cabling congestion. You just need to run a large number of cables.

Sorry.

the decision to switch from trays as the form factor, one of these as a tray, just just switching to racks as the form factor. That's a product decision. It's there wasn't a substantial technical barrier there.

Switching from from the, like, 64 to to 500 or so, there's a bit of Jensen math there, but there is at least a a genuine four x increase, which is coming from a much more complicated and difficult rack design. And so that that is actually, like, new new physical design to run more cables.

Yeah. I mean, it's let's sort of zoom in on this and look at the the wire density. I'll draw this diagram just once more so we have a bit of a cleaner version to work with, I mean, a larger version.

Let's say have some switches in the middle. Yep. And let's say I'm gonna have initially, I'm gonna start with just two GPUs on each side or two two trays of GPUs on each side.

And let's say maybe each tray wants to have two cables coming out of it. So I get some kind of I I physically run vertical cables that look like this running onto the switches. Now if I wanna double the number of GPUs in a rack, I need to run, like, literally twice the density of cables.

So I need to run, yeah, these as well.

Seems like there's a lot of space within a rack. I don't know. Just like the cables are, like, really big and Yeah.

So there is space outside the rack. Inside the rack like, these racks are, like I mean, as they become more optimized, these racks are very tight. So Yeah.

There's connector density going from from from the tray into the rack and the rack's backplane, and then the backplane itself has a has a really high density. Alright. There are other physical constraints, including, like, bend radius of cables.

Like, you don't wanna snap them and so on. Yep. Okay.

So it's literally the physical space to put a cable Yeah. That's constraining it. Yeah.

I had no idea. Interesting.

how about the rack is so big, and they're just, like we can't just stuff more cables in there. Yeah. So, I mean, rack design is not my expertise.

But, like, when I talk to to folks and what are the constraints they're up against, it's it's a combination of so what are the big physical things you're optimizing for? Space, weight of the rack. Like, it's actually really heavy, and so, like, you need enough metal top to not sag and fall.

But then you add more metal and it's heavier, and then power and cooling.

Yeah. Deploying in in larger scale up domains is a huge unlock. Yeah.

I mean, I've drawn here the sort of NVIDIA Blackwell deployment.

was a successful or it just seems like Gemini has that successful pre train for longer than some of the other labs. Not having been there at the time, I'm not sure how much is coming from, like, successfully deploying higher sparsity ratios, which which could be. It could also be I mean, there's a whole bunch of actual modeling things of, like, specifically, how do you do the mixture of experts?

We've seen the DeepSeek like, the DeepSeek mixture of experts has said, actually, activate more experts, but finer grained experts was a big innovation. I'm sure that there are many other innovations on the model architecture as well as on the training data. It's kind of hard to disentangle all of them.

But what shows up in terms of the limits of what you can do, The the active parameters, as we saw, is limited by the compute cost, and then the total parameters is limited by the scale up size.

So okay. So to answer that question, we're gonna need to talk about the communication patterns that so we've talked about the mixture of expert communication pattern. That is this all to all.

There's all to all. All to all. All to all very strongly favors full connectivity, which what we've kind of just shown here, and it favors being within one rack.

There are other kinds of parallelism besides expert parallelism, which which we just showed here. In the literature is tensor parallelism. This is with the trend towards smaller experts, this has become much less relevant, so we can ignore that.

But the other two things that we have available are data parallelism and pipeline parallelism. And they are actually much they can be a much better fit for using multiple racks. So let's focus on pipeline parallelism specifically.

This is one layer of MOE. I'm gonna have, like, a 100 more layers up above. I could decide at this point, for example, to move to a different rack change rack.

Now is that gonna become a communication bottleneck? So we can actually just solve for when this becomes a communication bottleneck. But before we do that algebraically, like, let's just sort of visualize it out and and sketch the path.

So we're gonna have a bunch this is another MOE layer, and we're gonna have another MOE layer here, and so on. So let's say I change rack here, and then some number of layers later, I change rack here as well. So our our our methodology that we're gonna use to determine whether we have a communication bottleneck in this, like, in this point where we change RAC is we're gonna compare the this this is the scale out scale out bandwidth requirements to the scale up bandwidth requirements.

Mhmm. So let's try this. And and, I mean, the the hint is gonna be that there's a lot more transcends here.

Like, we're sending many things here, whereas we're only sending one thing here. And then we're also maybe doing it many times.

Yeah. So number of activated GPUs. Right?

So, like, I I don't send to this GPU at all. Right? So there's an explosion from one to, like, three times larger here in in this diagram.

Yeah. The key thing is that I I didn't even need to send to this GPU at all, and so that's a big saving. I see.

Yeah. Okay. So we're gonna talk through sort of how much more what is the slowdown of to what extent is scale up a bottleneck over scale over scale out?

So we will directly jump to the ratio of the time spent on scale up, time on scale up, over the time spent on scale out. This is this is the quantity we're talking about. And the first consideration is that the scale up is, like scale up is is eight times faster than scale out, generally.

And so at a baseline, if the bandwidths were the same, we would have this one one over eight, which is coming from bandwidth bandwidth. But then we have some amount of expansion in in in how much data we're sending. So if one token comes in here Mhmm.

Then this one token gets routed to in the DeepSea case, it'll get routed to maybe 32 experts or or 16 experts. It gets routed to some number of experts. So this is the number of activated experts, number of activated experts.

And then it also this this same thing applies on multiple different layers. So maybe I'm gonna run two layers. So there's also multiple times number of layers per stage.

And just need to multiply the whole thing by two for the For the the optimal. Yeah. Yes.

Yes. And there's a factor of two. Thank you.

So what we would like is the for the scale up time to be greater than the scale out time, because, like, the scale up time is the more important and precious resource. And so we just we want this one we would like this number to be greater than or equal to one. And this really doesn't seem hard.

Like, we've we've there's just a factor of h that we need to overcome, so we need the product of these three things to be bigger than eight. Typically, we have a fairly large number of activated experts. It could be eight by itself.

And then we can increase the number of layers per stage a lot until until we satisfy this. I see. So what this ends up looking like is that I can, in fact, have an entire pipeline of racks where one rack does one layer, and then I move on to the next rack, and I do another layer, and then I move on to the next rack.

I can do another layer.

Isn't that I feel like that's interesting that the physical and The the model architecture matches like, the the the cutting matches the model architecture. Yeah. Exactly.

Yeah. I mean, it could have been something wackier with tensor parallelism and whatever. Yeah.

So, I mean, I think a way to think of it is I mean, okay. The Galaxy Brain way to think of it is, like, what are all the different dimensions in which a model is scaled up? And so there is it is scaled up by layers.

It is scaled up by the, like, demodel dimension. It is scaled up by the DFF dimension. It is scaled up by the number of experts.

Every single one of those numbers, you can choose to cut along. And if those numbers are big enough, it eventually becomes profitable to cut along there. Yeah.

And we have selected two of them. The other two, in the way typicals models are typically sized, are not profitable.

And so these complexities are real. Pipelining is a massive hassle. But it does give you some benefits.

And then you can then decide whether those benefits are worth the costs. Biggest benefit that shows up so it has some benefits in inference, maybe bigger benefits in training. In inference, what are we saving on?

Are we saving on memory time or compute time? Not really. We're just moving the memory time from one chip to another chip, or one rack to a different rack.

There's no actual benefit in runtime. However, what we are saving on is that the memory capacity is the amount of memory used per rack, if we think that the memory in a rack is a bottleneck, then there's a constraint on how sparse we can go. It pipelining allows us to massively reduce that bottleneck.

Okay. So so why do we do what what is this microbatching that shows up in shows up in pipeline parallelism? So the I'll focus on inference first.

It's it's a slightly simpler problem. So and I'm gonna draw so this is time, and then this is which rack rack we're on. And so the idea is that maybe I'll have, like, four racks.

So I've got an inference that is gonna, like, step through these four racks in some time like this. So, great. So this is inference number zero.

It runs at a certain batch size, and it steps through all all the pipeline stages like this. Now, if we were to say, well, we're gonna run inference number one here, like, this is clearly, like, a massive waste. Right?

Like, Three quarters of the time, each of the racks is doing nothing. So we don't actually run inference one here. We run it as soon as we can, which is immediately after inference zero finishes like this.

So and then we keep going. So if we hadn't filled this in, we would call this the pipeline bubble. When I've drawn it in this inference context where we're only going in a forwards pass, it's, like, obvious.

Like, why would you do this stupid thing? But in a training context, it's maybe less obvious. But in the inference context, it's sort of really natural to make this change.

Oh, interesting.

you need to Should we draw the the training diagram? Yeah. Yeah.

Yeah. Let's do that. Let's do that.

So so this is the inference diagram, and I'll call this four just so we don't have the wrong thing showing up there. So do the same thing for training now. We've got a forwards pass, but at some stage, we're gonna have to transition to a backwards pass.

So we'll we'll do some number of batches in the forwards pass. And then we're gonna transition to the backwards pass for everyone all in one go.

Yeah. I mean, smaller is always better, actually, is a way to put it. But from a ML convergence rate perspective, smaller is always better because basically you're getting the freshest information from the gradient descent.

But total trading time perspective. Total trading time perspective. It's worth like, smaller is worse from a systems perspective.

And so the optimum is the trade off between those two. Yeah. So you pick a batch size, and you and and then, like, for that batch size, you you do some amount forwards and then some amount backwards.

Yep. You asked why is there even a hard stop there? Pipeline parallelism, because of the fact that you've got this idle time here, which is the bubble, there are so many techniques in the literature for how to lay this out differently and and avoid that.

There are more complicated schemes called, like, zero bubble or one forward, one backward, which sort of interleave the forwards and the backwards in complicated ways.

More usefully, you can do the weight gradient step, but you can also weight it. Yeah. So, in inference, actually, the effect of pipelining on anything you care about, like batch size or latency, actually is neutral.

It it doesn't improve it. It doesn't make it worse. So if you look at the latency of this inference running it if it were pipelined versus if it were all on one rack, if it were all on one rack, we would just, like, slide all of the the boxes down and still put them in a row, and the latency would be the same.

So pipelining is neither better nor worse for latency, but it it does mean that you just use less memory per per rack, like memory capacity. Because now instead of needing the whole model, you only need a quarter of the model in Makes 10 sense.

So so even in inference, in fact, it is not used a ton. It say it reduces your memory capacity requirements. There's actually a huge surplus.

Like, I think you're saying that a rack of Blackwell has many, many terabytes, maybe tens of terabytes of that's much bigger than, like, a trillion parameter model. A trillion parameter model only needs one terabyte. And so it already fits, in fact.

And so there's not a huge benefit from pipelining because you're reducing a number that's already pretty small. But it does say that, theoretically, maybe you had too much memory. And maybe you could have done a different like, build a different hardware that has less memory, in fact.

If you were designing your hardware, like and you said, I actually didn't need that much memory because I don't need the weights to fit in one rack. I can fit the weights in eight racks. Then I could have maybe built a hardware that didn't have so much HBM per GPU.

this is Like, why why is Jetson shoving all this memory into these racks if Yeah. So you don't need it? Yeah.

So we've like, in in the equations we had here before we raised them, we were doing memory time, so memory bandwidth, then compute bandwidth. Let's now start looking at, memory capacity. Yeah.

So we'll start off with just, like, memory capacity without even thinking about parallelism scheme. And so the, like, the capacity of memory the demand on memory is the number of total parameters plus so so this is what we need to fit the weights in some system that we are using. Mhmm.

And then we need to fit the KVs as well. So KVs go as batch size times the length of the context times times the bytes bytes per bytes per check. Okay.

So what I was arguing about in this context and the case I was making for pipelining is that we will actually there are some techniques that allow us to solve this. Are there techniques that allow us to solve this? So let's let's consider So we're gonna run this on some number of GPUs, and and we're gonna say we're gonna have one extender, which is e is gonna be the expert parallelism.

So how many when we had this charting of expert layer across many GPUs, how much of that to what extent do we do that? How many GPUs? So we're gonna say that this is, example, 64.

And then p is going to be the extent of pipeline Pipelining. And so this is the number of racks, which, who knows, maybe maybe we'll we'll pick four or something like that. What we want to calculate so this is the this is like the total total memory rec requirement across the system.

But now I'm going to calculate a a memory requirement per GPU. So per per GPU memory requirement. We're gonna have I guess I'll use a lowercase c mem.

And, well, obviously, we just take all of these numbers and divide it by e and p. Really easy. So it's this n total plus the batch times length of context times bytes per toque.

All of this is divided by e and speed. Okay. So this is like, why is this correct as divided this way?

Well, we're we're we're saying we we knew that the parameters were perfectly divided amongst all the the GPUs in a rack. There are also the layers are perfectly divided amongst the the the different racks. So that works here.

And somehow, we're going to arrange I'll hand wave exactly how. Somehow, we can arrange the same perfect sharding of of the contexts across GPUs in a rack and and and and then based on layer across racks.

Yeah. For example. Yeah.

Yeah. So this is the place where we actually need to go back and analyze this batch size b. And you're making this comment that there's micro batching versus global batching.

So let's come back to this pipelining diagram here. We've got one batch going forward here. And then as I drew it, it kind of just, like, disappeared.

That's not really correct. If you think about how decode is working, I have a bunch of tokens that I have generated already. I do one forwards pass where I generate a new token.

And then and then I push like, then I write that to my KV cache, and then I do another forwards pass that generates the next token. So I'm actually gonna be running this batch zero in a loop. So, in fact, I go forwards.

Once I finish, I can start the next iteration of the loop up here. Yeah. So we'll just fill this in.

We'll we'll have a Nice. Yes. Yeah.

So we've got the two or three. Two or three. Two or three.

So let's split this batch. This batch will be the global batch size. So b is going to be the number of number of micro batches times the batch of like, the batch size per micro batch.

So how many micro batches do we need? So the number of micro batches in this diagram is four, zero one two three. And then the batch size per like, the the micro batch size, this is still this, like, 2,000 ish number.

This is the one that is, like Mhmm. This is the, like, 2,000 times sparsity. Sorry.

No. This is the 300 times sparsity. 300 times sparsity.

So the global batch size is the number of micro batches times the local batch size. Local batch size is set by this hardware parameter. The number of micro batches well, the number of micro batches is as small as possible such that we can wrap around and not leave any idle time when we wrap around.

So, if we had fewer, we would have had this idle time when we wrap around. And so, you can sort of just visually see that it is equal to the number of pipeline stages. Mean, a sort of proof by visual here, it is four, and it's four this way as well.

For sure, during massive scale training, this is done. It can be done for inference. I'm actually gonna make the case for why it is less attractive.

It is useful for weights, but not so useful for KBs. Okay. Yeah.

Yeah. The big challenge is so let's let's fill this in. The micro batch size here ends up being equal to the number of pipeline stages.

Yep. When we go back and substitute all of that into here, we get a number of pipeline stages times this little b showing up in here. And then when we factor this out, I'm gonna split this into, like, this plus into two two terms.

We get the full division by e times p over here. We still have division by e times p over here, but the p's cancel, this p and this p. They cancel.

And so what we find, if you increase the number of pipeline stages, the memory footprint for the number of weights keeps going down and down and down. Of course. But the memory footprint for the number of activations stays constant.

So so it it it doesn't actually work.

So why would it not help to be pipelining across multiple layers? Because then you don't have to store Yeah. You only need to store, like, one layer rather than two layers of KVs.

Right? Yeah. So so it helps from that perspective.

You're right. What's competing with that, though, is that you need to be keeping all of the racks usefully busy at a time. And so the number of sequences that are in flight simultaneously has gone up.

Yeah. Yeah. Makes sense.

Makes sense. Makes sense. So those exactly cancel, and and you end up not getting a saving per GPU.

Right. This is going back fundamentally to the point of you're you're not able to amortize across KV caches. Yeah.

Well, so first, we did you can amortize KV caches across batch size. Yeah. And now we're saying you also can't shard it across pipeline stages.

It sucks from both of those points of view. Yeah. Interesting.

Okay. So then what is that during inference? So, I mean, a, like, DeepSeg paper reports what they do, which is, like, they just do a lot of expert parallelism.

You should in effect, you should increase your expert parallelism up to your scale up domain size, and then do very little pipelining. Maybe none at all, maybe too just enough to make the weight storage not not too big of an issue. Those are the only two parallelisms that really make sense.

In the past, there was tensor parallelism, which was make cutting up within an expert, but the experts are so small now that that that is not a profitable optimization.

when they're doing inference, are just basically within a single scale up? Yes. Yeah.

I mean, you can look at how it depends on model size. Like, you could have a very large model, like one that exceeds the memory of a rack. And and and there, you should be doing a bit of pipelining.

Actually, I I mean the so pipelining doesn't help with context length. It totally helps with model size. And so because of the ability to do pipelining, at least a rack should not be a constraint on your ability to fit the model parameters.

I guess the other consideration you're asking, like, why hasn't it scaled up more, and why did bigger scale up domains help? So we we talked through one aspect of that, which is we kind of said it is not because of memory capacity. We have a solution to the memory capacity, at least with respect to model size.

Yeah. Interesting. Not with respect to KV cache size, but at least with respect to model size, we have a solution to memory capacity.

The other issue that shows up is latency.

what is the latency cost per per hop? This is very much dependent on the hardware. Yeah.

It's I would I can't say with a lot of authority. I think it's probably on the order of a few milliseconds, but it could be off by an order of magnitude there. Is four a realistic number of how many pipelining stages you might have?

Yeah. Yeah. Okay.

So that's that's not the case.

That's right. Yeah.

Or I don't know how many you said, but Yeah. Yeah. Ten milliseconds per token is actually a lot.

Yeah. If it if it goes from 20 to 30, right, or something like that. Yeah.

Yeah. And this is so, like, just to to chart the path that it goes through, here, you're going from your from your GPU or TPU or whatever, to a network card, which then goes to, like, a top of rack switch and then hops over to the other rack and does the same same thing in reverse. So you sort of have to sum up the latencies of these different things.

Sorry. This is the same thing as the DC This is the data switch? The data it may, in fact, go up to a decent switch and back.

Depends on deployment configuration. Got it. Yeah.

of the of this hop. There's also just the same TMEM latency, the memory time latency, is actually substantially, massively improved by larger scale up domains. So I'll recall TMEM down here.

TMEM for the weights, TMEM of weights, this was equal to the number of total parameters divided by the memory bandwidth. Which memory bandwidth are we talking about here? Is it just one GPU?

Or it's it's it's, in fact, is the number of GPUs that I can use in parallel to load these weights. So, I can't use different pipeline stages in parallel because they're not running at the same time. But I can use all the GPUs in my scale up domain in parallel to load the weights.

And so, this is actually extremely effective. So, basically, I end up with a term here. This this memory bandwidth term itself is equal to, like, scale up size.

And so this term doesn't increase a lot. It maybe increases 1.5 or two x per generation, but this one increased by, like, a factor of eight from these from Hubbard.

So so the reason the bigger scale of matter is not the memory capacity of the whole scale scale up, but really the memory bandwidth. Yeah. Yeah.

Pipelining totally solves the capacity problem, but but scale up size helps solve the bandwidth problem.

as a first thing. Like, if I just do a very sparse model and it's on, like, a little, like, h 100 box Yeah. The latency will will be really high.

Yeah. Okay.

This is a place where we have to do a bit of guesswork because, like, the the updated scaling laws and and the use and the model traffics are not reported, and so we have to guess there. But one way to look at it let me first just make a sort of a general heuristic claim. If I am if I had some, like, cost, and I've got a total cost, which is a sum of, like, cost a and cost b, like, maybe this is the training cost and this is the inference cost.

Yep. And so I wanna minimize this sum. For many for many curves that end up being the case, the minimum tends to be where these are where the costs are equalized.

That's something of a heuristic claim, but you can you can it tends like, there are many examples where it's true, like, where one is one over x and the other one is is x, for example. They tend to be minimized at at at the point where, they equal each other. It's also true for, like, e to the x and, like, e to the minus x and all kinds of other things.

Like so, basically, I've got some I've got some curve that's going down, some other curve that's going up, and they tend to be minimized up at this equal point. Heuristically, I will conjecture that that is true for the setup you described as well. Actually, showing that that would be true would require looking at the scaling laws and and, like, fitting these, like, weird exponents.

But but things that do follow power laws tend to tend to have this property. So I'll just make that claim and move on. So we're gonna say that the cost of training plus the cost of inference, we wanna equalize these.

We'll do pretraining only first because it's a little well, actually, we we can do all of it in general. So so actually, we'll we'll cost it as cost of pretraining. So number of number of number of active params times the data on pretraining.

So that's the cost of pretraining. There's a factor of six out here, which is the number of flops. There's the famous six n d formula.

And then in in RL, we have approximately the same thing. We've got, like, the same number of active parameters, but now it's the amount of data is the RL data. There is this extra, like, efficiency multiplier, which is or inefficiency, like the the inefficiency.

that.

if you're doing a backward pass on every single generation in RL, it would be six n d. Yeah. So this could be a smaller number.

Right? Like, this could be somewhere so They would at least be two. Yeah.

Somewhere in the range of two to six. So I'll just, like well, say somewhere into the range of two to six, something like that. Yeah.

And then and then we can add in the inference cost. The inference cost is two number of active times the data in inference.

And then inference is is just true. Yeah. So, we're gonna solve for, essentially, it may be a quality of all three of these terms, but that is ballpark where people are gonna be.

Like, labs have more information on on what is productive in doing more RL, for example, than versus doing more pretraining. I don't have that information. But I think a good ballpark is thirty thirty like, 33% split between each of them.

Actually, I'm not sure I understand the intuition for that.

and inference would be 50%. Yeah. That that's also a valid answer as well.

The because this is heuristic, I can't really argue for one versus the other. They don't differ by that much. Like, 33 versus 25 is only a small So, factor let's pick one of them.

All equal seems simple enough. And so we're just gonna solve for equality of them. It's pretty straightforward.

We can immediately see that the number of activated parameters totally disappears. And so let's factor that out. And we're gonna just say that data in pretraining.

I decided to do it your way. It's it's a little bit nicer, actually. So data in pretraining plus this oh, I didn't have the inefficiency over here either.

Inefficiency data in pretraining plus some multiple of, like, alpha times the data in RL is just gonna be end end end up equal to the sum sum of beta times the data in inference. So and then let's just, like, roughly size the alpha. This this this alpha, it's gonna be this is like the it's maybe somewhere in the range of two to six, two to six over six from this term compared to this term.

And then we've got an inefficiency term, which I would say is maybe in the range of, like, 30% or something like that. So so so this alpha is gonna be something like one and ten. One over 10, let's say.

And this beta here is is actually the same. It's it's a third it's one third times 33 to 30%. So it's also equals one in 10.

Something like that.

Yeah. Okay. We can make this, like, two in 10.

Uh-huh. Make it a bit bigger. Yep.

So, yeah, like, just write it out once more. Like, this is two and two over 10. This is one over 10.

So the number of inference tokens you have and this is just a function of, like, I've got hundreds of millions of tokens per second times my model is deployed for, I don't know, two months before I shift to the next version, that should determine the the number of tokens in in RL and pre pretraining, and then I guess we didn't do the equivalence between pretraining and and RL, so we'll do that here. Data in pretraining should be equal to, like, two over 10 times that make an RL for them to be cost equivalent. So sorry.

This one over I got it backwards. Like, we pay more cost when it's inefficient, so it's this needs to be one over. This tracing this back back forward.

Projects. Yeah. This this thing ends up actually being as written here, it's like oh, yeah.

So this is like 1.5, and and this is one.

think like if you do it with a spreadsheet, and like actually model it out, you might notice when the money's going down the drain. Yeah, yeah. So, yeah.

So I think this yeah. All of these end up being close in as modeled here. This 30% may have been a little bit too generous.

So let's say something like 1.5 here, and and leave this as a one here. So I think it like, at this point, you can almost read it off.

Like, the number of inference tokens should be about the same as the number of pretraining tokens should be about the same as the number of RL tokens within, like, factors that we're we're not able to reason about.

in terms of machine time. And so you if you're trying to equalize the RL and and pretraining time, then then you should have fewer tokens in order to have the same wall time.

equalizing in terms of data. I I I mean, I think starting with equalizing and cost is right, but Yeah. Depending on how you model the cost, this Yeah.

Comes close to equalizing in data.

if you think that people's power of prediction is not perfect Yeah. And and also you run the risk that you're that you make a model that is not a frontier model, and then you just throw it away. Then that kind of changes the cost trade off because there's some probability that applies to the inference, and you should de rate the inference tokens by some amount.

Right.

model? So I think we just have to make some real world assumptions here in order to do that. So so the inference tokens, we should totally be able to count.

Right? Like, so let's say a few 100,000,000. I don't know.

Maybe it's, like, 500,000,000 a second now. I don't really know. 500,000,000 a second times a model is deployed for two months before it becomes obsolete.

I don't really know. I can't do this in my head. Can you type it into a computer?

Okay. 2.6 times 10 to the fifteenth.

Okay. This number is probably too large. This because this is gonna be multiple models in a family.

We so let's Ah. Let's make it, like, five times smaller or 10 times smaller or something like that. Okay.

So we're estimating maybe 50,000,000 tokens per second per per specific model. The model is live for two months, and so this comes out to around 200,000,000,000,000 tokens. And then we want to compare that to active parameters on a Frontier model.

Active params? Sorry. Oh, so sorry.

Sorry. I meant tokens. Trained on a 150,000,000,000,000 tokens.

Interesting. Which is which is similar. Yeah.

Yeah. That's actually similar. So so data on pretraining.

This is not well cited, but You want me to not remove this side? It's not okay. And I think often active params number of active params could be in the range of, like, 100,000,000,000, something like that.

Yeah. Maybe maybe a little larger. So I'm assuming active frems of about a 100,000,000,000, and so multiply by 20 to get the Chinchilla token count.

So Chinchilla, d Chinchilla, would be around 2,000,000,000,000. And, yeah, and we see, like, we're about a 100 times larger than than that. Actually, what what does dChinchilla actually mean?

Got it.

Yeah. These kinds of numbers. This is also I mean, this is why you should just, like, approximate everywhere because, like, there's so big error bars on this.

But, yeah, no, it's kind of, like, empowering to just, like, set a equal to b and figure it out. Yeah. Yeah.

That's super cool. Okay.

Yeah. So, I mean, why specifically 50%? Let's let's sort of so so the high level, even in the first place, is there is some amount of increasing cost with with context links.

Yeah. And and we can bring that back up. That was the memory time versus the compute time.

So okay. So we we've put up these same occasions from before of the the time for memory fetches, which is the weights and and the k b cache, and then the the time for the compute, which is just the matrix multiplications for the weights. I will also draw the cost curve.

But this time, I'll do it as a function of context length instead of as a function of batch size. So this is time over yeah. Just just time.

And so this is the cost curve as a function of context length. We'll draw the compute. The the cost of the compute is actually constant as a function of context length.

There's no dependence here on context length. In reality, there is some dependence, but it is very mild dependence, so we'll ignore it. So this is the time for the compute.

This one. And then we'll also draw the dependence of the memory fetch on on context length. And this starts at a large number for the weights and then grows gradually with with the context length.

So maybe here, and then grow gradually with context length. And so you take the maximum, and you see there is this inflection point here. So now so this is the costs that that that, for example, Gemini might be paying.

And then you think, how how how might you put a pricing structure on top of that? You would like to ensure that no matter what the context length is, you are you are still profitable.

And so we've got a two tier pricing structure. Maybe we've got something that looks like this up to some next context. Fascinating.

So I think it says something about, given that the bump is at 200 k, it probably means that this is somewhat aligned with this crossover point, maybe not exactly aligned with. Fascinating. So we can actually probably even complete that calculation just to see where it lands out.

We can solve for the number of bytes per token, if we sort of make some assumptions about the number of active parameters. So, solving for the number of bytes per token, we're gonna assume, like, the point where we equalize the time of memory and the time of compute is at, let's say, 200 k tokens. So, we equalize these two.

We're also gonna just assume that the batch size is large enough that the memory time spent on weights is negligible, so we'll forget about this, and we'll focus on the actual memory time spent on KBCache. So that ends up saying copying this term over batch times lan context times bytes per token over MEM bandwidth is gonna be equal to number of activated params over our flops. And then we're gonna solve for bytes per token.

Match size was missing here. Shows up here. And then it cancels out by the time we get to here.

And and I I dropped the LENContext. So we can plug in numbers. This number this is this is this well, the reciprocal of the number that we saw before.

It's yeah. This is, one over 300, which is reasonably stable across many different hardware platforms. We conjecturally said that maybe number of activated tokens is, like, a 100,000,000,000.

And length of context we said was 200 k. Something is wrong here, though. Length of the context should be on the denominator, not the numerator.

So you said around two kilobytes. So let's just do a sanity check for this, for what this could be. There are two mechanisms that people do attention with a small number of bytes per token.

One is dense attention with a lot of reuse across layers. So Character AI has a blog post talking about that, alternating long and short context. And, like, in the Character AI kind of model, which also showed up in the Gemma models, the global context, which is really what we're talking about here, global context, was shared across all the layers.

And so to get this two kilobytes, you could get that, for example, as a d head of a 128 is is typical. And then, like, the number of bytes is typically number of attention layers times two times d head times number of queue heads. So this is the number of unique contexts per layer.

Do you do you share the the context across many layers, or do you do you it only once? So in Character AI like models, this number is one. We said this is a hundred hundred and twenty eight.

And this is a choice which typically ranges from one sorry. This is k v heads, I meant. So there was written a head, and a kvHead is that The kvHeads are the heads that are stored in memory, like, store the contents of the previous tokens.

The qHeads are the the retrieval heads there. They're only used temporarily, and they're they're used by the attending token. So in this autoregressive context Yep.

I've got k v heads associated with all of the context Yep. And then q heads associated with this new token here. But but but this head, the one twenty eight?

Oh, this is it it this this number is actually the same for oh, so this d head is the dimension of the vector. I I see. I see.

Yeah. Yeah. And number of k v heads is typically in the range of one to eight.

Yep. So, like, it is totally plausible to get this by, for for example, having eight k v heads and and a d head of a 128. That gives you exactly this number.

Yep. Or or you could have, like, fewer kV heads but more layers. Interesting.

Yeah. So that this is one way to get there via dense attention. There's also a way to get there via sparse attention where you increase all of these numbers, but then you have, like, a line over sparsity term.

So, yeah. I mean, I I think this number is plausible, if if maybe a little bit small. It's funny that they would leak so much information to their API pricing.

I mean, you are incentivized to price close to your costs because otherwise someone could scoop you. Maybe we can learn something about the difference in input versus output prices Yeah.

of the pass in general. Yeah. If we say, like, if we can think of decode as being a pass with one and then prefill being a pass with many.

Okay. Yeah. Yeah.

So maybe, like, prefix? Sure. Whatever.

I think maybe maybe sort of less draw actually how prefilled shows up here, if I may clarify. So we do a bit of decode like this. Yeah.

We may actually come back and do more prefill. Like like, if you think this is a chat session, the user says something, the AI generates response, and then the user says something else, then we prefill this. So, like, maybe this is the more common like, this is the general case rather than this.

And and and in fact, this is like you read a file or something. Read a file or just like the AI is responding to user input or a tool call or anything that's not Yeah. Exactly.

Generated. Yep. Okay.

Okay. So suppose we're here.

it would- Yeah. It's it's basically temporary. It it it doesn't even go to main memory.

Just ignore it. Okay. So then it would just be everything that came before.

is it not just that then? Yeah. There's actually no adjustment at all to the memory Great.

So so, yeah, there there is the time for one pass, but, actually, the amount of tokens is that that much larger. So I guess we want the cost per token, in fact, or the time per token.

Well, actually, for processing the entire batch. So in this, like, at this cost, we have processed this many tokens, like, LEN LENDER prefilled. Yeah.

Or I guess, pre yeah. Like, the of the pass. Yeah.

Yeah. No. Not not this prefix, but it's this cost.

Okay. Let's let's just do this in pass.

So the the result we want to work towards is that prefill is compute limited and decode is memory bandwidth limited.

Yep. Yeah. Like t on t on the y axis.

T we want the cost per token. So it'll be t over some stuff. T over length of the pass.

Mhmm. Yeah. That that'll be right.

then why is it cheaper? Cost higher? Yeah.

Yeah. So, I mean, we're gonna it's this division by length pass that that actually makes it all so okay. Yeah.

This is gonna divide out this is gonna divide out, but then we're gonna get a div all of this is gonna divide by length of pass, and it's gonna make the memory cost cheaper.

length length of the pass when it's one, that is decode. When it is bigger, that is pre fill. Okay.

I see. I see. I see.

That makes sense. Okay. Getting back to it.

Right.

Yeah. That's fine.

tokens, or the KV cache for a token. You can just produce it from scratch by computing it from the underlying token IDs, which are tiny. Or you can previously have produced it and stored it in a memory somewhere.

So the cost ratio is really talking about the ratio between those two mechanisms of producing it. A cache miss means you deleted it from all your memories, and you have to recompute it from the tokens directly. In fact, you can maybe even take that a step further and think about which memory tier do you store it in.

So you could store it in HBM. There are other slower and cheaper memories than HBM, like DDR on your host or Flash as well. And so one of the things you can do is a calculation of where it makes sense to be in each memory tier.

And this is related to how long you're going to store for. So we wanted to look at the cost of storage in a few different memory tiers and also the cost of rematerialization. So remat means the cost to rebuild all of the KBCh from scratch after you deleted it.

So we rematerialize it. And so, basically, this is gonna cost the length of the context. Actually, we'll look at cost per token so that we don't need to carry around this length of context everywhere.

So to rematerialize one token of KBChash, I just need to run I need to run a forward pass on the whole model. And and then so this is gonna be the compute time. I have to rerun the compute at whatever speed my GPU does it, and then I multiply it by my, like, GPU dollars per second.

Yeah. So there is a quadratic term in it shows up in the compute. As an approximation, I chose to remove it.

I'll just show you sort of quickly what that looks like. It's because so you have the if you look at the cost per token or the number of FLOPs per token, there is the FLOPs that are coming from doing the weight matrix multipliers Yep. As a function of context length.

And then there is the number of multipliers that comes from doing the KV cache, which is which goes up linearly with the the amount of stuff you attend to. The slope on this is so low that, like, when you when you draw it like this, it's, like, it's very well approximated by a flat line.

if this is true? Yeah. So there are two costs of long context.

One is the memory bandwidth cost, which we've spent a lot of time analyzing. That's this thing. And then the other one is the compute cost.

The compute cost is almost always, and sort of actually forced by fundamental principles to be a much smaller slope than than the memory bandwidth cost.

Yes. Token context links. I mean, sparse attention gives you a get out for sure because you get this square root, like, it gives But you a big I think it's like, if you look at the history of context lengths of models, from earlier models like GPT-three, maybe GPT-four, don't remember when the transition happened exactly, they shot up from about eight k to 100 k to 100 k.

And then for the last year or two, they've all been hovering around there. Yeah. I think that actually indicates that that that is sort of the reasonably balanced cost point, and going massively beyond that would be cost prohibitive.

So I actually don't see a very good path to solving that. Like, the memory the HBM is where it is is at where it is. It's not getting hugely better.

And and why doesn't attention solve it? Sparse attention is a big improvement. Maybe that is priced in already, perhaps.

It's not an infinite improvement because if you go too sparse, you lose too much quality. Yep. But, yeah, I mean, the empirical result is that the context links haven't been increasing that much.

And and I think it's because there is no solution to the memory wall here. Interesting. Like, so going too sparse just means, like, you're returning to a very small subset of the tokens, and the quality will get worse.

Yeah. So what is the cost of of these different ways of producing resynthesizing the KV cache? Computing it from scratch is based on my GPU time.

I have to do a certain amount of multiplies in order to of GPU time that I spend in order to produce it. Store in HBM. This really goes as my I think I had a number here, which was the bytes per token.

So I need to I need to just have some number of bytes per token, And then I need to store this in the HBM. So it's gonna use up some of my HBM capacity. So a way to think of this is that, like, if I have too many of these things sitting in my HBM, like, if I fill up my HBM with just KV caches that I'm not using, I can't use that GPU.

And so how do I price that? Maybe I say that the cost of it is proportional to the fraction of the HBM I'm using. So so there's also times GPU dollars.

And then let's just do one more memory tier and say something like DDR store in DDR instead. The same kind of thing goes up for Flash and and for DDR. I put these in the wrong columns, actually.

I meant to make two columns. The the distinction I want to make is that there is the time to cost to retrieve. And then there's a cost cost to store, costs to hold hold on.

Yeah. And so this is like there's a cost per second, whereas this is like an instantaneous cost. So rematerialization has a cost to retrieve and has zero cost to store it because we deleted it.

This is the one that I put in the wrong location. This is this is actually the cost just to hold on, so I will rewrite it. Okay.

So we have this is the like, if we're just storing it in HBM, it has this sort of cost profile. And then if we store in DDR, it's actually gonna take some time. So, we get the same thing here, bytes bytes per token over DDR capacity times DDR cost a second.

But but now this has a a cost to retrieve that is is higher than the HPM, because we need to copy it into the HPM. And so this is bytes per token over DDR bandwidth bandwidth, and then this consumes some amount of the DDR as well. And every scale up has DDR and Flash?

This is really a deployment question, and so you you can choose that. NVIDIA does deploy in this form. It has it has both.

Why isn't the cost to retrieve HBM the memory bandwidth or the bytes divided by memory bandwidth? Yeah. I mean, it depends what what you define a retrieve to be.

Here, I'm defining retrieve to be move it into HBM so that you can start actually doing inference on it. And so, like, sort of by definition. And because if it's already in HBM, you can be doing compute while you're getting it from HBM Deshp, for example.

Okay. Yeah. So so these are three things, and I I guess I ordered them wrong.

In general, if you're if you're balancing two costs and you've got different memory different tiers in the memory hierarchy, you should expect as as this cost goes up, this cost should go down. So you can kind of see where the zeros are, and, like, I should have ordered them. This one first, this one second, and this one third.

So if you're going to hold on to it for for a very short amount of time Yeah. Then all of this is, like, multiplied by the hold time. Yep.

This one is and so is this one.

Yeah. Right. Which suggests that the five minutes is HBM and the hour is DDR.

I I think that's, like, I think that's a pretty good assumption. It could if you look at the numbers, it might also turn out that it's one tier down, and it's TDR versus Flash is the Yeah. Okay.

Interesting. And the price difference, I think, was I'll look it up.

tokens. Base, which means remap? Yeah.

That's 5. This is 5? Woah.

$5.

So actually, we might actually be able to determine the which memory tier it is by by the durations, actually. The duration probably tells it to actually And the the six five minutes versus one hour. Yeah.

Exactly. I think this will probably end up being it's gonna be the drain time of the memory tier that you're in. And so what that means is, like, given that I'm I know that I'm gonna be holding something for five minutes, I would like to have pick a memory that I can read every five minutes.

Like, I can read the whole memory once per five minutes, ballpark. So that is the drain time of the memory. So if I take the the the, like, all the storage storage capacity over storage bandwidth I would like this to be, like, equal to five minutes or something like that.

Mhmm. And so, actually, we did this calculation for HBM. For HBM, we know that this number is twenty milliseconds.

So HBM is much too short, like, much too small. DDR could be about an order of magnitude or or too off from this, and so this is probably in the order of, like actually, I think it might even be in the in the seconds, like, one to ten seconds. And then this is really I don't have these numbers memorized, but generally, as you go to slower tiers, Flash is plausibly in the order of one minute.

And then, like, spinning disk, which is massively different, think, is on the order of one hour. So this might actually identify that the tiers are probably flash and spinning disk. Sorry.

Why why is this the calculations, the storage cap divided by the bandwidth? So you you've got a bunch of different memory tiers. Like, we've listed four of them.

Your choice of which memory tier is you want to minimize the cost. So what fraction of the device are you using? You're using some fraction of the device for the holding onto it, and then you're using some fraction of the device to retrieve it.

And so let's say I'm using, like, 10% of the device. And I would equalize those two fractions. That's a sign that I've hit the right thing.

So let's say I've got some runtime here. Like, I I'm gonna hold on for all of this time, and then so this is the time hold. And then there's gonna be some amount of time here, which is time retrieve.

Mhmm. And I want I mean, basically, to equalize the costs, these two costs, I want the retrieval time to be equal to the hold time times the, like, fraction of capacity.

Because, like, this is the the retrieval time. Yeah. I mean, this is Yeah.

Yeah. This is how many other things I can hold simultaneously.

get all your things in there and out. Yeah. Basically.

Yeah. I I think that probably indicates that this is the the two tiers of Flash and and spinning disk. I'm kinda shocked to see spinning disk being used at all because it's such an old technology.

Yeah. But yeah. I mean, it's also crazy that it's so slow that it takes an hour to load its full capacity into it.

Then Like, it's a really unattractive technology, but it's useful in some places. Yeah. So we're sitting down because I wanna ask you some questions that I guess don't need a platform.

So, I mean, the, like, the mixing, like I I try to look for other examples where mixing like, scrambling mixing shows up as well. There's actually almost even, like, a physical example where, like, you're stirring something, or you're making a cake, you wanna stir the batter. And, like literally, the idea, like, first stir it this way and then stir it this way is, like, actually not too bad of an approach.

But beyond that, like, in back to the digital world, there are some differences, and the one you call out is a pretty strong difference. The way it shows up, like, what makes neural nets like, if you just randomly initialize a neural network, actually, maybe it's a reasonable cryptograph like, a cipher as well because, like, the random initialization is gonna jumble stuff in a complicated way. It may even, like, do what you want.

Who knows? The thing that makes it interpretable is the gradient descent. So you can differentiate a neural network and get a meaningful derivative.

And we do a lot of work to not overcomplicate the derivative. So the residual connection keeps it, like, contained and simple, and and so does, like, the layer norm stuff that we do. One of the biggest attacks against cryptographic ciphers is also to differentiate the cipher.

Ciphers run-in a different number field. They run-in the field of two elements, so just binary, whereas neural nets run, like, in theory in the field of real numbers. And so you have to differentiate with respect to, like, binary numbers.

But you can absolutely differentiate a cipher, and this is called differential cryptanalysis. And, like, basically, what it says is that if you take a small difference of the input, how, like, it it's quite difficult to make the difference of the output be small. Like, oh, like, the whole job of a of a well designed cipher is to make the difference in output very large.

So I I guess the distinction is that the the optimization goals at that point are about complexifying.

Yeah. Yeah. Test.

Yeah. Yeah. I mean, in fact, this is, like, this is actually a place where you get exactly the sort of avalanche property that ciphers have as well.

Like, adversarial attacks on typically, like, image classification models, right, are, can I find a perturbation of the image that a very, very small perturbation of the image that totally changes the classification totally changes the output? That is the common case in ciphers, whereas that's the, like, undesired case in in neural nets for sure. Yeah.

Okay.

well, in general, cryptography, like, creating a new cipher is a very, very dangerous proposition. Like, almost all of them are broken. Like, 99% of them are broken.

So probably a bad place to start. But the other direction has been very like, in in at least one very clear case, quite productive. So there's this construction sorry, a construction that exists in ciphers and then was imported into neural nets called a Feistel cipher, Feistel network.

So the idea is that you may have some function f, which is not invertible. But you like the function because it, like, does interesting things. Like, it it it does an MLP, for example Yeah.

Or or it mixes in an interesting way. You'd like to build something out of this that is invertible. So the construction we're gonna make is gonna actually be a two input function rather than a one input function.

And we're gonna apply f of x. We need to actually remember what x was, so we're gonna stick x over here so that we can work backwards. And then we also can't drop y.

So we're gonna remember y, and we're gonna add them together. And so we form this tuple. So the the way to invert this, like, if you think I have this output and I want to recover x and y, well, I can easily recover x.

That's right there. I just read it off. And then to recover y, I like, if this thing was called z, I can I can recover y by z minus f of x, because I've already recovered x?

So so that means that this construction is invertible. This was used in ciphers, like, a ton. Still is used.

It's one of the main mechanisms of constructing ciphers. Often, want ciphers to be invertible, especially the layers of ciphers you want to be invertible, because that has better heterographic properties. This has actually been ported over into neural nets.

There's a twenty seventeen, eighteen paper called REVNets, reversible networks. And what it does is it actually makes the entire like, you can apply it to any network, like a transformer network. You can make I do a forwards pass, but then I can actually run the entire pass backwards as well.

So the whole neural network is invertible with exactly this construction. And so this paper reversible networks, like, applied to some layer, like a transformer layer, for example, we've got this function f, which is our transformer layer. Now, normally, we would have just an input and then a residual connection coming out, and it gets added like this over here.

Mhmm. But now the variation of this is going to be we've got two inputs, x and y. So we've got x and y inputs.

X goes through the function, gets added to y, And then this becomes the new x, the output x. And then this x becomes the output y. So really, what this is doing, this is, like, this is actually sort of doing if you think of two layers back, this is actually the thing you mentioned before.

It's actually doing the residual connection from two layers back. Like, this y came from the previous layer and was the residual connection there. But because of this construction, the whole thing is invertible.

Why do I care? What does invertible matter for? The big thing that it can be interesting for is for training.

If I think of a forward pass of training, so I will let's say I have four layers. I run them in zero, one, two, three order. I have to write all of the activations to HBM.

And so I get an HBM footprint here that is kind of like linear linear in number of layers. Yep. So this this actually can be the largest memory footprint during training.

And so this is normal training. And then and then I run the backwards pass, and I read it kind of in reverse. Like, I I run them in sort of forward pass goes forward, backward pass goes backwards, and I have to read them back out.

The idea of this RevNets paper is that because it's invertible, I don't need to store this at all. I can completely rematerialize it when I'm running my backwards pass. So I I run my forwards pass, and then when I'm running my backwards pass, I'm simultaneously in lockstep, undoing all of the forwards pass steps that I did in order to to have the activations that I need here.

So this ends up being a memory saving, which is a nice idea. Interesting.

Spending more memory to save compute is generally profitable given where Yeah. Hardwares are today. Interesting.