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

22:26So let's say I've picked a batch size that I'm going to run at. Maybe I pick, you know, this batch size. And so like, well, and by the way, this intersection point is the same intersection point here. So I pick this batch size. I know that it's going to take, for example, maybe it's something like 20 milliseconds is a common place to sense up landing. What I'm going to produce is, like, so this is a timeline of what is running on the GPU. It's going to start a new batch every 20 milliseconds, regardless. And so, sorry, this is 20, this is 40.

23:05You can think of this as a schedule for the train. A new train departs every 20 milliseconds. Any passengers who are ready, board the train. If the train is full, then they wait to the next train. if the train is not full, the train's going to go anyway. And so in terms of what that means for queuing latency, it means that the worst case is that a request arrives just after the train departed. It has to wait for the next train, so that's up to 20 milliseconds, and then it has to wait for that train to complete. And so the worst case latency is 40 milliseconds. Sure, how is a 20 milliseconds derived?

23:39I 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 kb's. 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 the chip. And so that is capacity divided by bandwidth. with that tends to be 20 milliseconds on many different generations of HPM. The units make sense.

24:17You would have a byte divided by bytes per second. Yeah. So for example, I mean, on, I think the Rubin generation, it is something like 288 gigabytes divided by 20 terabytes per second.

24:36And this looks like it comes out to about 15 milliseconds. let me just make sure i understand what it's saying i mean i understand how why the units can't yeah the sort of unit analysis but what is it saying is

24:52we can evacuate and replace the hbm in this amount of time and so we don't want to be in a situation where the hbm is not big enough that we're not you know actually able to keep write everything you want to it or take everything out of it or we don't want to be in a situation where our ability to write back and forth is so big or so small compared. Yeah, there's sort of two scenarios. Why don't we pick a latency that is bigger than 15 milliseconds? And 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.

25:31It's like almost all reads because the weight matrices are read only and then almost all of the KV cache accesses are reads. So in, like, let's say I run 30 milliseconds, I can read all of HBM twice. But what's the point of that? Like, I don't want to read the weight matrices twice. I don't want to read the KVs twice. Yeah, makes sense. Makes a ton of sense. Okay, so a couple of actually quick questions. One, if it is the case that the optimal batch size is something like 2000, and that actually true, it's totally dependent on sparsity, it's not dependent on the model size or anything. I mean, sparsity shows up in model size, but beyond that, it only depends on sparsity, not on scale.

26:07That's a very interesting result. And that seems to imply that you can... One question is, how much of a 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. I don't know, is 2 ,000 users at the same time a lot? It doesn't seem like a lot? We can do a bit of analysis on this, which would be actually... 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. So what does this batch size mean in terms of tokens per second of the system?

26:38So 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 every time in tools, which is, let's say, which is, this thing is equal to the 15 milliseconds, 20 milliseconds number. So this ends up being batch size itself times

27:04So, like 64 times b. And so this ends up being around 2000 times 64. So like 128k token-specific. So this is sort of in more digestible units. Like, it's hard to reason about concurrent users, but what is the global traffic for a system?

27:29When you look at some of the announcements, sometimes the API providers will brag about how much traffic they have. 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. So about a thousandth of that range. Gemini, yeah. But I mean, Gemini is big. So that's actually 1 ,000th of Gemini is a lot. To actually be competitive at scale, you need to be able to serve at least 1 ,000th of Gemini. That's interesting. Cool. Okay, so the more sparsity you have, the less compute you need. And it does seem that as batch sizes get bigger, compute ends up being the bottleneck, according to this analysis.

28:20So then the question is, how far can you take sparsity? That is to say, as the sparsity ratio increases, as you have fewer and fewer active parameters relative to total parameters, how much is performance of the model degrading? And is it degrading faster than your saving compute by increasing the sparsity factor? Yeah, so performance, equality of the model, rather than speed of the model. 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. All right. Should we follow the paper now or so?

28:58Yeah. 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 techniques have changed a lot. DeepSeq mixture of experts was a big change in how it worked. There have been older papers which are G-shared switch transformer. So the actual empirical results are going to depend on all of that.

29:34But 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.3b dense across, you end up seeing that, for example, in this case, the 64 expert 370 million activated parameters model is as good as a dense 1.3 billion model. So in some sense, it's actually not amazing returns where you need to increase total parameters 100-fold to get the equivalent of 10x as many active parameters.

30:13Yeah, I mean, actually, even more so. Yeah, it's a huge increase in parameter count for a modest increase in. Yeah. So in this case, actually, it's what is it? 4x? 64x for 4x. Yeah. So while it is true, I guess, that you get this benefit of being able to economize on your compute time if you increase sparsity, naively it would seem like, oh, that's a trade-off worth making. but if this you're decreasing this by 2x and then having this go up by 8x every time you double sparsity. So 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.

31:01And so just keep running a larger batch size. From the point of view of the analysis we've done here, this is pure win. Keep doing it. Keep doing it until you run out of available users, basically. 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 a reasonable tradeoff. The other tradeoff that shows up here is that it also consumes memory capacity, which we've only reasoned about memory bound with it, but it also consumes memory capacity. So let me just make you understood. You're saying we want to spend less time computing, therefore we do more sparsity.

31:49To make that work, we need bigger batch sizes, which means we need more memory capacity. to have more sparsely. Yeah, so maybe this would be a good point to actually talk about how a mixture of experts layer is typically laid out on a rack of GPUs or something. Yeah, yeah, makes sense. Yeah, where were we? Sparse mixture of experts. Yes. Maybe how we lay that out on a GPU. Yep. So let's zoom in on the mixture of experts layer first and sort of draw what that looks like. So, we typically will have some kind of a router layer, which is making the decision of where we route the tokens to. So, we have tokens coming in here, they go through a router layer, and then we have a bunch of different experts.

32:43I'll draw a few more to line some up. And then the router will make a decision, which experts am I going to 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 each expert itself is a normal MLP. It has a up projection and then a down projection and a non-linearity 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.

33:31The token is also passed through here and it gets added to the result of the MOE layer. So this is a normal MOE layer. What I want to talk through is how this is mapped to a GPU rack and what this means for communication. because I think this will start to show some of the limits of how sparse we can go. So the standard practice here, and 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 DeepSeq 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.

34:18This is not a power of two. So we'll just simplify and say we're only going to use 64 of them. Just ignore the other eight. It's not a big deal. And so we have four experts per GPU. Very simple. For the sake of the diagram, I'll actually just say, let's say we have two experts per GPU. So 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, some tokens stored centrally here. They get routed to all of these experts. And so there is some communication cost paid here. There's the same communication cost paid on the output.

35:03And 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 their model. So this is an all-to-all traffic pattern. So when you say any GPU in the pretense, the router is more than one GPU? Yeah, so I drew this as one router. In reality, you would actually have many copies of the router. And so you would have as many routers as GPUs, in fact. As the incoming traffic. Yeah. So these are 64 GPUs. These are 64 GPUs. It's actually the same GPUs.

35:47We just draw them as separate because they're serving different purposes. So at this point, any GPU can be sending to any GPU. So this all-to-all pattern of communication that shows up and how the blackwall racks are configured is a perfect fit for 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.

36:32And so the rack-to-rack communication ends up being a substantial bottleneck. So the fundamental thing here is that one rack actually 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. It may be worth you explaining what exactly a rack is, the differences in bandwidth between a rack and within a rack, and the all-to-all versus not all-to-all nature of communication within versus outside. Yeah, 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.

37:13So generally, a rack is a... It is a physical structure. 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. What constrains it being a certain size is power delivery, weight, and cooling ability. It ends up being about this size in many cases because of these physical constraints. So then when I deploy a data center, a data center may have thousands of these racks. So I've got one of these tall racks, it's got a bunch of GPUs in it, and so on. And then I put another rack next up. You make it sound so easy.

38:01Yeah, right? I just like drop them in. In NVIDIA's case, the communication topology is actually, they put the GPUs 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. And then they run a bunch of cables. every single GPU has cables going to the switches in the middle.

38:39So every GPU goes to the 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 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. 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 some data center switch.

39:22So data center switch.

39:27And then all of the GPUs will have some connectivity up to some data center switch somewhere. But this is about times, like this is the scale out.

39:40And it tends to be about eight times slower 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 going to be wanting to talk to the GPUs here. And so, like, just on average, like, when I look at where the tokens on these GPUs want to go, half of the tokens want to go inside the rack, that's great. They can use the fast scale up network, but half the tokens are going to want to leave the rack and go to the other rack. And that's not as good. They're going to need to use a much slower network. And so that becomes the bottleneck on the all-to-all pattern.

40:24A different choice would be, well, why don't I have a big switch here and sort of like, and connect everything to some big switching, like a much bigger switch that actually combines the two racks together. So 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, is that question you just asked, basically, why isn't it a bigger scale-up? Yeah, exactly. Why not just have a million chips and scale-up? What has changed that has allowed NVIDIA to go from Hopper was 8, then Blackwell is 72, and now Rubin will be, is it 500 or something?

41:09Yeah, 500 or something, yeah. What has allowed that to happen? From Hopper to Blackwell is mostly just the decision to switch from trays as the form factor. One of these is a tray, just switching to racks as the form factor. That's a product decision. There wasn't a substantial technical barrier there. Yeah. Switching from the 64 to 500 or so, there's a bit of Jensen math there, but there is at least a genuine 4x increase, which is coming from a much more complicated and difficult rack design. So that is actually like new physical design to run more cables. And the cable complication is just the cost of figuring out which cable hops to which signal.

42:00Yeah, I mean, let's sort of zoom in on this and look at 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 I have some switches in the middle. Yep. And let's say I'm going to have, initially I'm going to start with just two GPUs on each side, or 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 physically run vertical cables that look like this running into the switches. Now if I want to double the number of GPUs in a rack, I need to run like literally twice the density of cables.

42:43So I need to run these as well.

42:51It's really not a question, but if you look at a physical data center, it seems like there's a lot of space within a rack. I don't know, just like the cables are really big. Yeah, so there is space outside the rack. Inside the rack, these racks are, I mean, as they become more optimized, these racks are very tight. So there's connector density going from the tray into the rack and the rack's backplane. And then the backplane itself has a really high density. There are other physical constraints, including like bend radius of cables. You don't want to snap them and so on. It's literally the physical space to put a cable that's constraining it.

43:31I had no idea. Interesting. That seems surprising. that like, 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 folks and what are the constraints they're up against, 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 to not sag and fall, but then you add more metal and it's heavier and then power and cooling. And so all of those are competing for, like, modern racks are pushing all of those to very extreme physical limits.

44:09Deep work is by its nature quite aversive, so even things which seem like work, like Slack and email, can be easy ways to distract yourself. So I often wish that I could just turn the internet off. But if I'm prepping for an interview, even if I have the papers and books on hand, it's still super useful to be able to do a back and forth in the LLM so I can break down concepts and research follow-ups. Google's new Gemma 4 is the first open model that allows me to have this kind of fully disconnected focus machine. It's small enough to run on my laptop, but good enough to actually be useful. So to prep for this episode, I downloaded Reiner's scaling book and shut off the internet.

44:42I was able to have Gemma help me understand the material and answer my questions. If you want an LLM that you can run locally on your laptop or even your phone, you should check out Gemma 4.

44:54When was GPT-4 released again? It was 2022 or 2023? Three, three, yeah. And it was rumored to be over one trillion parameters. And it seems like only now and within the last six months have models been getting released that are significantly more parameters than a model released three years ago. when supposedly there should have been this scaling in the meantime, is the reason that we were just waiting for racks with enough memory to hold a 5 trillion parameter model along with its KV cash for enough users for a lot of sequences or if you're doing RL, kind of a similar consideration of actually holding the KV cash for all the batch of problems you're trying to solve.

45:40So if you look at like Hopper, you had eight Hoppers, and I think that's 640 gigabytes as of 2022. With Blackwell finally, which was deployed, what, 2020? Very recently, maybe last year. Last year? Yeah. You finally have a scale-up with on the order of like 10, 20 terabytes, which is enough for like a 5T model plus KB cache. Yeah. Deploying in larger scale-up domains is a huge unlock. Yeah. I mean, I've drawn here the sort of NVIDIA Blackwell deployment. The Google deployment has actually had very large-scale domains for a long time. And that also explains why Gemini seemed to be ahead. Like, was Gemini 2.5, was a successful, or it just seems like Gemini has that successful pre-train for longer than some of the other labs?

46:24Not having been there at the time, I'm not sure how much is coming from, like, successfully deploying higher sparsity ratios, 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 DeepSeq 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, So the active parameters, as we saw, is limited by the compute cost.

47:07And then the total parameters are limited by the scale-up size. Yep. When you're operating within a single scale-up domain, is that a consideration specifically for either forward or backward? Or specifically for pre-fill versus decode? or is it preferred to always be within a scale-up? Yeah. Whatever kind of workload you have, whether you're doing a pre-training run or whether you're doing RLL generation or whether you're doing inference for users. Yeah, really interesting. So, okay. So to answer that question, we're going to need to talk about the communication patterns. So we've talked about the mixture of expert communication pattern.

47:53That is this all-to-all. is all-to-all. All-to-all. All-to-all very strongly favors full connectivity, which is what we've kind of just shown here, and favors being within one rack. There are other kinds of parallelism besides expert parallelism, 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.

48:37So let's focus on pipeline parallelism specifically. This is one layer of MOE. I'm going to have like 100 more layers up above. I could decide at this point, for example, to move to a different change rack.

48:59Now, is that going to become a communication bottleneck? So we can actually just solve for when this becomes a communication bottleneck. But before we do that algebraically, let's just sort of visualize it out and sketch the path. So we're going to have a bunch, this is another MOE layer, and we're going to 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.

49:28So our methodology that we're going to use to determine whether we have a communication bottleneck in this point where we change Rack is we're going to compare the, this is the scale out bandwidth requirements to the scale up bandwidth requirements.

49:52So let's try this. I mean, the hint is going to be that there's a lot more sends 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. So that's going to be what makes the difference. Can I try to guess? Just out of curiosity to see if I'm actually understanding. It seems like you're sending like batch size into the rack. In here? Yes. but the communication within a rack is sort of batch size times number of GPUs. Yeah, so number of activated GPUs, right? So, like, I don't send to this GPU at all, right? So there's an explosion from one to, like, three times larger here in this diagram.

50:38Yeah. The key thing is that 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 going to talk through sort of how much more, what is the slowdown of, to what extent is scale up a bottleneck 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 the quantity we're talking about. And the first consideration is that the scale up is like... Scale up is eight times faster than scale out generally. And so at a baseline, if the bandwidths were the same, we would have this one over eight, which is coming from bandwidth.

51:34Bandwidth. but then we have some amount of expansion in how much data we're sending. So if one token comes in here, then this one token gets routed to, in the DeepSea case, it'll get routed to maybe 32 experts or 16 experts, gets routed to some number of experts. So this is the number of activated experts, number of activated experts.

52:08And then it also, the same thing applies on multiple different layers. So maybe I'm going to run two layers. So there's also multiple times number of layers per stage. And there's a need to multiply the whole thing by two for the, um, for the, yes, yes. And there's a factor of two. Thank you. Um, so what we would like is the, for the scale up time to be greater than the scale out time, um, 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. Um, and this really doesn't seem hard.

52:53Like we've, we've, there's just a factor of eight 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 we satisfy this. 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. It's interesting to me that the best parallelism strategy in practice ends up being one which physically resembles the actual architecture.

53:32It's not some galaxy brain thing. It's like, oh, we have experts. We're going to put them on different GPUs. Oh, we have different layers. We're going to put them on different racks. I feel that's interesting that the physical and... The model architecture matches, like 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 demodeled dimension, it is scaled up by the DFF dimension, it is scaled up by the number of experts.

54:11Every 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. and we have selected two of them. The other two, in the way models are typically sized, are not profitable. So there's a talk by Ilya where he says, today we know not to do pipeline parallelism. And Horace gave my friends and me, I hate that it sounds like a Dr. Seuss quote. But he gave us a lecture on these different kinds of parallelisms and he said, the problem with pipeline parallelism is that it other than the bubbles, it creates these architectural constraints.

54:52Yes. Like Kimi, for example, has these residuals where attention attends to the... A few back or something. Yeah, layers a few back, and so that becomes hard to implement in this way. Yeah. And I guess we didn't really fully articulate even what is the benefit that we're getting from pipelining. Yeah. And so these complexities are real. Pipelining is a massive hassle, but it does give you some benefits.

55:23And then you can then decide whether those benefits are worth the costs. The 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.

56:04Pipelining allows us to massively reduce that bottleneck. I guess the opposite connotation to this, which actually, before this interview I was chatting with Axel, who's a GPU performance engineer at Jane Street, he was explaining, well, to do pipelining, you have to do micro-batches rather than full-batches. and if you do micro batches then you're by definition not able to amortize the weight loading the weights that's right across all the users or all the sequences and so the positive connotation of that is you don't have to use this memory the negative connotation is that of that is that we can't amortize loading the weights across all those users maybe it's worth explaining why you had to do micro batches because you yeah so we draw the pipeline bubble Yeah.

56:52Yeah. Okay. So why do we do... What is this micro-batching that shows up in pipeline parallelism? So I'll focus on inference first. It's a slightly simpler problem. And I'm going to draw... So this is time. And then this is which 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 going to step through these four racks in some time like this. So this is inference number zero. It runs at a certain batch size, and it steps through all the pipeline stages like this. Now, if we were to say, well, we're going to run inference number one here, this is clearly a massive waste, right?

57:45Like three quarters of the time, each of the racks is doing nothing. Yeah. 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. 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 obvious. 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. So this is sort of obvious, but the difference between micro-bash and bash doesn't matter at all in inference because you can just call whatever you want, whatever.

58:30Yeah. It only matters in training because there is an optimal batch size. Yes. and before you do the backwards step, you want to have accumulated, before you do a full backwards step, you want to have accumulated all the sequences in that batch and if you want to do pipeline and training, in order to avoid that bubble, you need to... Should we draw the training diagram? Yeah, let's do that. Let's do that. So this is the inference diagram and I'll call this forward so we don't have the wrong thing showing up there. So let's do the same thing for training now. We've got a forwards pass, but at some stage, you're going to have to transition to a backwards pass.

59:10So we'll do some number of batches in the forwards pass.

59:20And then we're going to transition to the backwards pass for every run-on-one one go.

59:33So the inference part is the same here, but then we do a hard stop at this point and then transition everyone to backwards pass. Similar numbering like this. It may be worth clarifying. The reason there is that hard stop is because you want to do a whole batch at once for the backwards step. And then there is an optimal size for how big that batch should be. Yeah, I mean, smaller is always better, actually, is a way to put it. But from an 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.

1:00:08Total trading time perspective, smaller is worse from a systems perspective. And so the optimum is the trade-off between these two. So you pick a batch size and then for that batch size, you do some amount forwards and then some amount backwards. you asked why is there even a hard stop? Pipeline parallelism, because of this, 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 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.

1:00:49You can mine Bitcoin in that. Yeah, right, right. More usefully, you can do the weight gradient step, but you can also make it correct. So in inference, actually, the effect of pipelining on anything you care about, like batch size or latency, actually is neutral. 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 slide all of 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 does mean that you just use less memory per rack, like memory capacity, because now instead of needing the whole model, you only need a quarter of the model.

1:01:36Makes a ton of sense. So basically, no brainer to use pipelining during inference, but there's this harder trade-off during training. So even in inference, in fact, it is not used a ton. 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. 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.

1:02:16But 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 HPM per GPU. Last week, Porosi was kind enough to give me and my friends a great lecture on large-scale pre-training systems. And there were some concepts that I wanted to animate for a write-up on my blog, like how weights shard and gradients flow depending on the parallelism that you're using.

1:02:54So I gave Cursor my lecture notes and a sketch that I made during the lecture, and I asked it to visualize a specific hierarchical collective that Horace had explained. The first version was already pretty good, and then I was able to use design mode to select and tweak any specific components from there. I was able to do all of this without a clear end state in mind. Cursor's Composer 2 Fast model was quick enough that I was able to iterate almost instantaneously. I could try an idea, test the results in the built-in browser, and immediately make any changes. I went through 10 different versions in under 20 minutes.

1:03:24If you want to check out this animation, I published it along with the lecture notes in a blog post. The link is in the description. And if you want to try out this kind of iterative design flow for yourself, go to cursor.com slash Larkash to get started. so macro question everybody's talking about the memory wall right now the memory's getting super expensive there's not enough memory smartphone volume will go down 30 because there's not enough memory hyperscalers are spending this is shocking if i'm dylan said they're spending 50 percent of their capex this year on memory on memory that's believable yeah so but it's so like what is hyperscaler capex that's like high hundreds of billions maybe a trillion and They're spending half of that on memory.

1:04:05Okay, so that is this huge constraint. That's why we're not going to get new laptops and phones this year. But at the same time, we have too much memory. Like people are willing to put too much memory into these systems. Right. So this is... Like why is Jet Set shoving all this memory into these racks if you don't need it? Yeah, so in the equations we had here before we raised them, we were doing memory time, so memory bandwidth and 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.

1:04:38And so the capacity of memory or the demand on memory is the number of total parameters plus, so this is what we need to fit the weights in some system that we are using. And then we need to fit the kVs as well. So kb's go is batch size times the length of the context times the bytes per token. 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 consider, so we're going to run this on some number of GPUs and we're going to say, we're going to have one extent which is E is going to be the expert parallelism.

1:05:41So 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 going to say that this is fact, for 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 we'll pick four or something like that. What we want to calculate, so this is the, this is like the total memory requirement across the system. But now I'm going to calculate a memory requirement per GPU. So per GPU memory requirement. We're gonna have, I guess I'll use a lowercase C mem.

1:06:35And well, obviously we just take all these numbers and divide it by ENP. Really easy. So it's this n total plus the batch times length of context times bytes of toque. All of this is divided by e times p. Okay, so this is like, why is this correct divided this way? Well, we're saying we knew that the parameters were perfectly divided amongst all the GPUs in a rack. The layers are perfectly divided amongst 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 the contexts across GPUs and a rack, and then based on layer across racks.

1:07:27And so four is the number of racks? Yeah, for example.

1:07:35So 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 um let's come back to this pipelining diagram here um 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 um 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 kbcache and then I do another forwards pass that generates the next token.

1:08:13So I'm actually going to 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.

1:08:26So we'll just fill this in. We'll have the...

1:08:34Oh, uh, yeah. So we've got the two or three, two or three. Uh, so let's split this batch. This batch will be the global batch size. So B is going to be the number of micro batches times the batch size. 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 4, 0, 1, 2, 3. And then the batch size per, like, the micro-batch size, this is still this, like, 2 ,000-ish number. This is the one that is, like, this is the, like, 2 ,000 times sparsity. Sorry, no, this is the 300 times sparsity. 300 times Smart City. This is how big the train that takes up every 20 milliseconds is.

1:09:34Right, yes. This is going to be the 20 milliseconds train. 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 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. I mean, proof by visual here, it is four and it's four this way as well. But you can sort of look and see that it goes along here and then it wraps around number of pipeline stages.

1:10:14Yeah, very basic question. This is what is actually done? Mm-hmm. Okay. As in Frontier model today, we'll actually have, during inference, have pipeline. For sure during massive scale training, this is done. It can be done for inference. I'm actually going to make the case for why it is less attractive. It is useful for weights, but not so useful for kb's. The big challenge is, so let's fill this in. The micro batch size here ends up being equal to the number of pipeline stages. When we go back and substitute this, that's all of that into here, we get a number of pipeline stages times this little b showing up in here.

1:11:08And then when we factor this out, I'm going to split this into like this plus into two terms.

1:11:17we 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.

1:11:29They 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, but the memory footprint for the number of activations stays constant. So it doesn't actually work. Like if most of your memory ends up like, once you do enough pipelining and it's really not much, like even two is often enough, this term becomes very small. This becomes the dominant term. The KV cache becomes the dominant term. Yeah. I know this is wrong. I'm just trying to think out why my train of logic here is wrong. If you have many different, you're pipelining through many different stages.

1:12:10The KV values are not shared between layers. 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 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... Yeah, yeah, yeah. Makes sense. Makes sense. Makes sense. So those exactly cancel. And you end up not getting a saving per GP. Right.

1:12:41This is going back fundamentally to the point of you're not able to amortize across KV caches. Well, so first we did, you can't amortize KV caches across batch size. And now we're saying you also can't shard it across pipeline stages. It sucks from both of those points. Yeah, yeah, yeah. Interesting. Okay, so then what is done during inference? So, I mean, DeepStick paper reports what they do, which is like, they just do a lot of XO 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 two, just enough to make the weight storage not too big of an issue.

1:13:25Those are the only two parallelisms that really make sense. In the past, there was tensor parallelism, which was cutting up within an expert, but the experts are so small now that that is not a profitable optimization. So this goes back to the question, does that mean that Frontier Labs, 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 there you should be doing a bit of pipelining. Maybe it's extremely sparse, for example, and that would be a reason to do it.

1:14:05So I guess this goes back to the question about, or this goes back to the promise at the beginning of the lecture, which was, this will actually tell you about AI progress as well. To the extent it is the case that model size scaling has been slow until recently because, let me make sure I understand the claim. The claim would not be, you could have trained across more racks. It was just that it would not have made sense before. Like we didn't have the ability to do inference for a bigger model easily. Actually, I made the clue. So pipelining doesn't help with context length. It totally helps with model size.

1:14:40And 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, why hasn't it scaled up more and why did bigger scale-up domains help? So we talked through one aspect of that, which is we kind of said it's 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.

1:15:16I was just about to ask, so what is the going from rack to rack? What is the latency cost per hop? This is very much dependent on the hardware.

1:15:28I 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. Is four a realistic number of how many pipelining stages you might have? Yeah, yeah. Okay, so that's not that much. On a small number of pipelining stages, this is not a huge latency impact. Wait, I guess it's 10 milliseconds per token. That's right. Two times four-ish, or I don't know how many you said, but 10 milliseconds per token is actually a lot. Yeah, if it goes from 20 to 30, right, or something like that. Yeah. So just to chart the path that it goes through, here you're going from your GPU or TPU or whatever to a network card, which then goes to a top of rack switch and then hops over to the other rack and does the same thing in reverse.

1:16:21So you sort of have to sum up the latencies of these different things. Sorry, this is the same thing as the DC switch? It may, in fact, go up to a DC switch and back. Depends on deployment configuration. Got it, yeah. And because it's decoded in sequential, it's also not the, like, they stack up across the stages. Yes, they do. You can't do them at the same time. That's right, yeah. Okay, so I guess this brings us back to the question then. Is the size of the scale-up at all relevant to why AI model sizes or whatever have been what they have been over the last few years, whether through training or through inference?

1:16:57Yeah, so, I mean, we talked about latency of the hop, of this hop. There is also just the same TMM latency, the memory time latency, is actually substantially, like massively improved by a larger scale of domains. So I'll recall TMM down here. Tmem for the weights, Tmem of weights.

1:17:27This was equal to the number of total parameters divided by the memory bandwidth.

1:17:38Which memory bandwidth are we talking about here? Is it just one GPU? or it's in fact, it 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 memory bandwidth term itself is equal to scale-up size. Times memory bandwidth per GPU. Yeah, times GPU bandwidth. And so this term doesn't increase a lot. It maybe increases 1.5 or 2x per generation, but this one increased by a factor of 8 from Hubbard.

1:18:26So the reason the bigger scale-up matter is not the memory capacity of the whole scale-up, but really the memory bandwidth. Yeah, yeah. Pipelining totally solves the capacity problem, but scale-up size helps solve the bandwidth problem. And the bandwidth problem helps you do longer context lengths, which is more and more relevant as these models get more authentic. Yeah, it lets you just run the model at lower latency as a first thing. If I just do a very sparse model and it's on like a little H100 box, the latency will be really high. Yeah. Okay, a super tangential question. there's chinchilla scaling which tells you how big should a model be relative the amount of data you're going to train it on um but now obviously you're not just trying to optimize for the highest quality model you could get with training compute you want the best results a user can get yes the mixture of training and inference compute so then there's a question of how much should you over train a model such that that compute amortized over training and inference is minimized to get a certain performance.

1:19:34But now with RL inference, there's, or RL, there's another consideration, which is you're going to do some minor pre-training. That pre-training will be used both for RL generation and then for inference for the final user. And by over-training here, I mean, while it would have been more efficient just from a training computer perspective to have a bigger model that you train for less time because it can learn faster. Maybe you get a smaller model. you spend more computerating it than you otherwise would have, but now it's cheaper to give it to users. Basically, okay, let me give a question more concrete.

1:20:06How much more than chinchilla optimal are models over-trained? Yeah. And has that changed as a result of RO generation? This is a place where we have to do a bit of guesswork because the updated scaling laws and the model traffics are not reported, and so we have to guess there. But one way to look at it...

1:20:29Let me first just make a sort of a general heuristic claim. If I have 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. Yeah. And so I want to minimize this sum. For many curves that tend up being the case, the minimum tends to be where the costs are equalized. That's something of a heuristic claim, but there are many examples where it's true, like where one is one over x and the other one is x, for example. They tend to be minimized at the point where they equal each other. It's also true for e to the x and e to the minus x and all kinds of other things.

1:21:20So basically, I've got some curve that's going down, some other curve that's going up, and they tend to be minimized at this equal point.

1:21:30Heuristically, 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 fitting these weird exponents. But things that do follow power laws tend to have this property. So I'll just make that claim and move on. So we're going to say that the cost of training plus the cost of inference, we want to equalize these.

1:22:03We'll do pre-training only first because it's a little, well, actually we can do all of it in general. So actually we'll cost it as. cost of pre-training, so number of active params times the data on pre-training. So that's the cost of pre-training. There's a factor of six out here, which is the number of flops. There's the famous 6ND formula. And then in RL, we have approximately the same thing. We've got like same number of active parameters, but now it's the amount of data is the RL data. There's this extra like efficiency multiplier, which is, or inefficiency, like the inefficiency.

1:22:50Which is the fact that you're not trading on all your rollouts. Well, yeah, there's that. And then the other perhaps even bigger inefficiency is that this involves a substantial amount of decode and often decode runs at less MFU than training. Okay, so if you're doing a backward pass on every single generation in RL, it would be 6ND. Yeah, so this could be a smaller number, right? This could be somewhere... It would at least be 2. Yeah, somewhere in the range of 2 to 6. So we'll just say somewhere in the range of 2 to 6.

1:23:27And then we can add in the inference cost. the inference cost is two, number of active times the data in inference. I think the way I said it was super garbled, just for the audience, maybe. Forward plus backwards per parameter is six. Forward alone is two. That's why RL, where you might, you're definitely going to generate all the trajectories, but you might or might not train all the trajectories is two to six. Yes. Yeah. Thank you. And then inference is just two. so we're going to solve for essentially maybe a quality of all three of these terms that is ballpark where people are going to be like labs have more information on what is productive in doing more RL for example versus doing more pre-training I don't have that information but I think a good ballpark is 30-30 like 33 % split between each of them actually I'm not sure I understand the intuition for that another naive model could have been that RL plus pre-training would be 50 % an inference would be 50%.

1:24:30Yeah, that's also a valid answer as well. 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 factor.

1:24:44So let's pick one of them. All equal seems simple enough. And so we're just going to 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 going to just say that data in pre-training... I decided to do it your way. It's a little bit nicer, actually. So data in pre-training plus this... Oh, I didn't have the inefficiency over here either. Inefficiency. Data in pre-training plus some multiple of alpha times the data in RL is just going to end up equal to the sum of beta times the data in inference.

1:25:37So, and then let's just like roughly size the alpha. This alpha, it's going to 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 this alpha is going to be something like 1 over 10.

1:26:09And this beta here is actually the same. It's a third. It's one third times 30%. So it's also equals 1 in 10, something like that. If both of them are 1 in 10, that kind of implies that there's never a backward pause on RL? Yeah, okay. We can make this like 2 in 10. Make it a bit bigger. So, yeah, like just write it out once more. Like this is 2 over 10. This is 1 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.

1:26:50That should determine the number of tokens in RL and pretraining. And then I guess we didn't do the equivalence between pretraining and RL, so we'll do that here. Data pretraining should be equal to like 2 over 10 times data in RL for them to be cost equivalent.

1:27:12So, 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.

1:27:23So this tracing this back forward, this, this thing ends up actually being as written here. It's like, yeah, so this is like 1.5 and this is one.

1:27:38Billions of dollars of the compute just flowed the other direction. Yeah, right. I 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 as modeled here. This 30 % may have been a little bit too generous. So let's say something like 1.5 here and leave this as a 1 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 pre-training tokens, should be about the same as the number of RL tokens within factors that we're not able to reason about.

1:28:15But then, so it looks, sorry I'm making a basic algebra mistake, it seems like there should be less RL tokens than pre-training tokens? Yeah, that's in general right, because RL is less efficient in terms of machine time. And so, if you're trying to equalize the RL and pre-training time, then you should have fewer tokens. in order to have the same wall time. This is quite interesting that I never thought about it in terms of how much equalizing in terms of data. I mean, I think starting with equalizing in cost is right, but depending on how you model the cost, this comes close to equalizing in data.

1:28:52That if every single user who uses, basically if you forge GPT to be trained optimally, every single user who uses GPT-5, the total amount of tokens that they stream should equal the total amount that have gone into pre-training. Yeah. And the total amount of tokens that I got in pre-training is the sum of all human knowledge. So like each model should generate the sum of human knowledge on the output that it gets on the input. Yeah. So, I mean, which way are people going to err? Like if you think that people's power of prediction is not perfect and also you run the risk that you make a model that is not a frontier model and then you just throw it away.

1:29:31then that kind of changes the cost trade-off because there's some probability that applies to the inference and you should derate the inference tokens by some amount. And then can we back out how much more compute than Chinchilla Optimal for a given-sized model? So I think we just have to make some real-world assumptions here in order to do that. So the inference tokens we should totally be able to catch, right? Like, so let's say a few hundred million, I don't know, maybe it's like 500 million tokens a second now. I don't really know. 500 million tokens a second times a model is deployed for two months before it becomes obsolete.

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

1:30:212.6 times 10 to the 15th. Okay, 2.6 times 10 to the 15th. Okay. This number is probably too large because this is going to be multiple models in a family. So let's make it like five times smaller or 10 times smaller or something like that. Okay, so we're estimating maybe 50 million tokens per second per specific model. The model is live for two months. And so this comes out to around 200 trillion tokens. And then we want to compare that to active parameters on a frontier model. I don't actually know the latest rumors, but some... Do you know? Somebody told me 150 trillion... Active params? Sorry, I meant tokens.

1:31:15Trained on 150 trillion tokens. Interesting. Which is similar. Yeah, that's actually similar. So data on pre-training. This is not well-cited, but... You want me to not remove that? No, it's fine. And I think often active params, number of active params, could be in the range of like 100 billion, something like that. Yeah. Maybe a bit larger. So I'm assuming active params of about 100 billion, and so multiply by 20 to get the chinchilla token count. So chinchilla, de chinchilla, would be around$2 trillion. And yeah, and we see like we're at 100 times larger than that. Actually, what does de chinchilla actually mean?

1:31:58Like the token count for pre-training for the chinchilla scaling law would recommend, I guess. Oh, I see. So how much is it over-trained? Got it. So yeah, like the ratio of this 200 trillion or 100 trillion parameters over the, like, the potential optimal of 2 trillion, that's the amount that's over-trained, which is like a factor of 100 over-trained. That's whatever. Okay, so if you consider this right here, to the extent this is in the right ballpark, just by thinking about, okay, you kind of want everything to be equal in terms of compute. Here's, if that OpenAI also realizes that in their serving are certain amount of tokens per second, that tells you how much data went into the pre-training of GBD5.

1:32:43Even if it's like 50 % off or something, that is sort of wild that you can sort of first principles for these kinds of numbers. This is why you should just approximate everywhere because there's so big error bars on this. It's kind of empowering to just set A equal to B and figure it out. Yeah, that's super cool. In the spirit of trying to deduce things, we can publicly look up the prices of the APIs of these models. And maybe you can learn something from that. So first, with a longer context, Gemini 3.1 is 50 % more expensive if you go over 200K tokens than if we're below 200K tokens. I mean, at a high level, I understand why that might that be, but why specifically 50 %?

1:33:34Yeah. Yeah. So, I mean, why specifically 50 %? Let's sort of... So, the high level, even in the first place, is there is some amount of increasing cost with context length. And we can bring that back up. That was the memory time versus the compute time. So, okay. So, we've put up these same equations from before of the time for memory fetches, which is the weights and the kb cache, and then the time for the compute, which is just the matrix multiplications for the weights. I will also draw the cost curve.

1:34:22But 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 time. So this is the cost curve as a function of context length. We'll draw the compute. 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.

1:34:56This one. And then we'll also draw the dependence of the memory fetch on context length. And this starts at a large number for the weights and then grows gradually with the context length. So maybe here and then grow gradually with context length.

1:35:18And so you take the maximum and you see there is this inflection point here. So now, so this is the costs that, for example, Gemini might be paying. And then you think, 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 still profitable. Interesting. And so we've got a two-tier pricing structure. Maybe we've got something that looks like this up to some context. Fascinating. So I think it says something about, given that the bump is at 200K, it probably means that this is somewhat aligned with this crossover point. Maybe not exactly aligned with.

1:35:55Fascinating. 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 going to assume like the point where we equalize the time of memory and the time of compute is at, let's say, 200K tokens. So we equalize these two. We're also going to 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 KB cache.

1:36:38So that ends up saying copying this term over batch times len context times bytes the token over mem bandwidth

1:36:53is going to be equal to number of activated params over flops.

1:37:03And then we're going to solve for bytes per token.

1:37:27Match size was missing here. shows up here and then it cancels out by the time we get to here.

1:37:39And I dropped the LEN context.

1:37:44So we can plug in numbers. This number, this is this, well, is the reciprocal of the number that we saw before? Yeah, this is like 1 over 300, which is reasonably stable across many different hardware platforms. We conjecturally said that maybe number of activated tokens is like 100 billion.

1:38:04And length of the context we said was 200k.

1:38:09Something is wrong here. The length of the context should be on the denominator, not the numerator.

1:38:291667. Like about one kilo, almost two kilobytes. That is plausible, actually. So you said around two kilobytes.

1:38:44So 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 dhead of 128. is typical.

1:39:24And then the number of bytes is typically number of attention layers times 2 times dhead times number of qheads. So So this is the number of unique contexts per layer. Do you share the context across many layers or do you use it only once? So in character AI-like models, this number is one. We said this is 128.

1:40:09And this is a choice which typically ranges from one. Sorry, this is KV heads, I meant. So there was written a head and a KV head is that? The KV heads are the heads that are stored in memory, like store the contents of the previous tokens. The Q heads are the retrieval heads there. They're only used temporarily and they're used by the attending token. So in this auto-aggressive context, I've got KV heads associated with all of the context. And then Q heads associated with this new token here. But this head, the 128. Oh, this is, this number is actually the same for, Oh, sorry. This d-head is the dimension of the vector.

1:40:49Ah, yeah. And number of kV heads is typically in the range of 1 to 8. Yeah. So, like, it is totally plausible to get this by, for example, having 8 kV heads and a d-head of 128. That gives you exactly this number. Yeah. Or you could have, like, fewer kV heads but more layers. Yeah. So this is one way to get there via denser tension. There's also a way to get there via sparse tension where you increase all of these numbers, but then you have like a lot of a sparsity term.

1:41:20So yeah, I mean, I think this number is plausible if maybe a little bit small. It's funny that they would leak so much information through 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. And what that tells us about decode versus pre-fill in these models. and I think last I checked it's like 50 % more expensive or something like that I don't remember what I've seen in the past is like 3 or 5 times more expensive that makes more sense let's say it's 5 times more expensive okay this is the compute to process the next token in decode suppose you're doing pre-fill where you're not just processing the most recent token you're processing all the tokens in parallel so I want to say that that it would be this times len, um, len pre-fill?

1:42:16Sure. Len's of a pass in general. Yeah. If we say like, if we can think of decode as being a pass with one and then pre-fill being a pass with many. Okay. Yeah, yeah. Um, so maybe like prefix? Sure. Whatever. Um, okay. Memory. So you're not storing the KV cache if you're for the tokens that are the pre-fill tokens. I think maybe sort of let's draw actually how pre-fill shows up here. If I may clarify, so we do a bit of decode like this. We may actually come back and do more pre-fill. Like if you think this is a chat session, the user says something, the AI generates response, and then the user says something else, and we pre-fill this.

1:42:57So maybe this is the more common, like this is the general case rather than this. 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 tool call or anything that's not exactly generated. Yeah, okay. Okay, so suppose we're here. So you will need to load.

1:43:19Basically, you will have calculated all of this previously. So just the KV of everything that came before.

1:43:28But what is the memory cost of this? well, memory bandwidth cost of this. If you're doing flash attention, it would... Yeah, it's basically temporary. It doesn't even go to main memory. Just ignore it. Okay, so then it would just be everything that came before. So is it not just that then? Yeah, there's actually no adjustment at all to the memory time. Okay, great. Oh, so it's a very trivial change to accommodate. So this term is making it 5x more expensive. Now, why would that be? Or what does that tell us about... What are we trying to learn here? What does that actually tell us? What variable does it help us clamp?

1:44:12Well, the compute has presumably gotten 5x. The only thing that could have changed is the compute is 5x more expensive as a result. So yeah, there's the time for one pass, but actually the amount of tokens is that much larger. So I guess we want the cost per token, in fact. or the time by token. So I'm not sure I understood. This is for processing the next token in prefix? Well, actually for processing the entire batch. So at this cost, we have processed this many tokens, like len, len, pre-fill. Yeah. Well, I guess pre, yeah, like of the pass. Yeah, not this prefix, but it's this cost. Okay, let's just need this in pass.

1:44:59So this is 5x more expensive? Input is 5x more expensive? No, output is more expensive. Output is 5x more expensive. So the result we want to work towards is that pre-fill is compute limited and decode is memory bandwidth limited. Why don't we do this? Why don't we just chart it with like len pass on the x-axis? Yep. And t on y-axis. T, we want the cost per token, so it'll be T over some stuff. T over length of the pass. Mm-hmm. Yeah, that'll be right.

1:45:42Okay, so...

1:45:56Okay. It gets me confused about this. Length pass is the... It seems like this should be higher when you're doing pre-fill. Pre-fill has a bigger length pass. Yeah. Right. But then why is it cheaper? Why is it cost higher? Yeah, yeah. So, I mean, we're gonna... It's this division by length pass 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... All of this is going to divide the length of pass, and it's going to make the memory cost cheaper. Okay, yeah, let me think about this then. Okay, so let's do one line for...

1:46:34Basically, we'll have four different lines. Let's do the... Let's do pre-fill first, and so... Actually, let's do decode first. Oh, so actually, I mean... 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. So t compute, if you have basically just this divided by length pass, it's just this amount. So this actually does not vary based on t. So it'll just be some flat value like this. And this is t compute. And then this is like, this is... That's decode. decode right um now tmem if you have this whole thing divided by len pass well it doesn't really matter what's up there it'll just be something that looks like this right yeah let's say this is tmem this is decode again so as the length of the prefix goes up or pass, your memory bandwidth time declines.

1:47:58And that means that to the extent that you were bottlenecked on memory bandwidth before, you can avoid being bottlenecked on memory bandwidth. The fact that they are charging 5x less for pre-fill than decode does suggest that they are bottlenecked on memory bandwidth to quite a degree, such that for them at least, because T is equivalent to cost, right? It's the cost of renting a compute. This is actually like, this would be at one and this would be at five. That's right, that's right. Yeah. So it is, in fact, tremendously memory bandwidth. The real graph looks something like, the real graph looks something like, like that.

1:48:40Yeah, I mean, still it crosses, but. Yeah, exactly. So yeah, let me do it this way. Yeah, that's right.

1:48:50and then the this is the gap on decode between the memory and the compute time yeah yeah okay interesting another interesting one would be why cash hits are so much cheaper yeah okay so i think if i remember correctly cash hits are like 10x it's more expensive to write to cash according to the pricing on all these models. But if you do hit a cache, it's 10x cheaper. So what is going on with... Presumably, this is the cost of keeping something in HBM rather than just evacuating it. But if you do keep it in HBM, then it's cheaper to load again? Right. So there's two ways you can produce tokens or the KV cache for a token.

1:49:40You 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've 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.

1:50:19And 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 want 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 KB cache from scratch after you deleted it. So we rematerialize it. And so basically, this is going to 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 KVCache, I just need to run a forward pass on the whole model.

1:51:13And then, so this is going to be the compute time. I have to rerun the compute at whatever speed my GPU does it. And then I multiply it by my GPU dollars per second. I was very extremely naive question. Why is there not a quadratic term? Yeah, so there is a quadratic term. It shows up in the compute.

1:51:44As 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 multiplies as a function of context lengths. And then there is the number of multiplies that comes from doing the KV cache, which goes up linearly with the amount of stuff you attend to. The slope on this is so low that when you draw it like this, it's very well approximated by a flat line. So you start to notice the effect of the quadratic or the linear term up in the millions of tokens or so.

1:52:30So just not super relevant. So what is the reason that there's no company which has over a million token context link? 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 the memory bandwidth cost. And so the primary thing that limits you to have really large contexts are memory bandwidth and memory capacity, which is exactly this effect.

1:53:09like um and so there's this idea that dario said on the podcast and others have said which is we don't need continual learning for agi in context learning is enough and if you believe that then you have to think that we had to get to 100 million token 100 million billion context length to have an employee that is the equivalent to working with you for a month now maybe that's no longer true as far as attention or something yeah but um yeah if you think that then as a some MLinfor thing would have to change to allow for 100 million, like the memory bandwidth to allow for 100 million token context lengths.

1:53:45I mean, sparse attention gives you a get out for sure because you get this square root, like it gives you a big improvement.

1:53:55But I think it's like, if you look at the history of context lengths of models,

1:54:04from like earlier models like GPT-3, maybe to GPT-4, I don't remember when the transition happened exactly. Like they shot up from like about 8K to 100K, 200K. And then for the last year or two, they've all been hovering around there. I think that actually indicates that that's sort of the reasonably balanced cost point and going massively beyond that would be cost prohibitive. Not because of the compute cost, but because of the memory bandwidth. cost yeah um so i actually don't see a very good path to solving that like the memory the hbm is where is it's at where it is uh it's not getting hugely better and why doesn't sparse attention solve it the sparse attention is a big improvement um uh maybe that is priced in already perhaps um uh it's not an infinite improvement because if you go too sparse you lose too much quality But yeah, I mean, the empirical result is that the context links haven't been increasing that much.

1:55:03And I think it's because there is no solution to the memory wall. Interesting. Like, so going too sparse just means like you're attending to a very small subset of the tokens and the quality will get worse. Yeah. So what is the cost of these different ways of producing,

1:55:21resynthesizing the kvcache? Competing it from scratch is based on my GPU time. I have to do a certain amount of multiplies of GPU time that I spend in order to produce it. Storing HBM.

1:55:42This really goes as my, I think I had a number here, which was the bytes per token. So I need to just have some number of bytes per token. and then I need to store this in the HBM. So it's going to use up some of my HBM capacity. So a way to think of this is that if I have too many of these things sitting in my HBM, 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 there's also times GPU dollars.

1:56:23And then let's just do one more memory tier and say something like DDR. Store in DDR instead.

1:56:33The same kind of thing goes up for flash and for DDR. I put these in the wrong columns, actually. I meant to make two columns. The distinction I want to make is that there is the cost to retrieve.

1:56:50and then there's a cost to store, cost to hold it, hold on. 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've deleted it. This is the one that I put in the wrong location. This is actually the cost to hold on. so I will rewrite it.

1:57:34Okay, 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 going to take some time. So it's like we get the same thing here, bytes per token over DDR capacity times DDR.

1:58:01cost, a second. But now this has a cost to retrieve that is higher than the HPM because we need to copy it into the HPM. And so this is writes by token over DDR bandwidth.

1:58:21And then this consumes some amount of the DDR as well. And every scale up has DDR and flash? There's really a deployment question, and so you can choose that. NVIDIA does deploy in this form. It has both. Why isn't the cost to retrieve HBM the memory bandwidth, or the bytes to write it by memory bandwidth? Yeah, I mean, it depends on 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, sort of by definition. Because if it's already in HBM, you can be doing compute while you're getting it from HBM to HBM?

1:58:54Yeah, for example. So these are three things, and I guess I ordered them wrong. In general, if you're balancing two costs and you've got different tiers in the memory hierarchy, you should expect as this cost goes up, this cost should go down. So you can kind of see where the zeros are, and I should have ordered them this one first, this one second, and this one third. So if you're going to hold onto it for a very short amount of time, then all of this is multiplied by the hold time. This one is, and so is this one.

1:59:38And interestingly, they have different prices to write for, and you specify this in the API, for five minutes versus an hour. Yeah, which suggests that the five minutes is HBM and the hour is DDR. 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 DDR versus flash. Yeah, okay, interesting. And the price difference, I think, was, I'll look it up. Okay, so the base input tokens is five per million tokens. Base, which means rebound. Yeah, that's five. This is five? Five hours. To like retrieve quote unquote and then the to write to

2:00:30presumably HBM write for five minutes is 6.25 So actually we might actually be able to determine the which memory tier it is by by the durations actually the durations probably tells it to actually Five minutes versus one hour. Yeah, exactly. I think this will probably end up being, it's going to be the drain time of the memory tier that you're in. And so what that means is like, given that I know I'm going to 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.

2:01:12So if I take the storage capacity over storage bandwidth, I would like this to be equal to five minutes or something like that. And so actually we did this calculation for HBM. For HBM, we know that this number is 20 milliseconds. So HBM is much too short, like much too small. DDR could be about an order of magnitude or two off from this. And so this is probably in the order of like, actually, I think it might even be in the seconds, like one to 10 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.

2:01:56And then like spinning disk, which is massively different, I think is on the order of one hour. So this might actually identify that the tiers are probably flash and spinning disk. Sorry, why is this the calculation? So the storage cap divided by the bandwidth? So you've got a bunch of different memory tiers, like we've listed four of them.

2:02:17Your choice of which memory tier is like, you want to minimize the cost. And so you are like, 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 want to 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'm going to hold on for all of this time and then so this is the time hold and then there's going to be some amount of time here which is time retrieve and i want i mean basically to equalize the costs these two costs um i want the retrieval time to be equal to the hold time uh times the like fraction of capacity

2:03:22um because like this is the the retrieval time uh yeah i mean this is yeah this is how many other of things I can hold simultaneously. Basically, just like, hey, you want to store things in there for so long such that the amount of time it's in there is kind of the time to get all your things in there and out. Yeah, basically. I think that probably indicates that this is the two tiers of Flash and spinning disk. I'm kind of shocked to see spinning disk being used at all because it's such an old technology. I mean, it's also crazy that it's so slow that it takes an hour to load its full capacity into it.

2:03:58It's a really unattractive technology, but it's useful in some places. So we're sitting down because I want to ask you some questions that I guess don't need to platform. You have this extremely interesting blog post where you talk about how at a high level, the architecture of different cryptographic protocols looks a lot like neural networks. And there's this conversion evolution where they both need to jumble information across all their inputs for cryptographic protocols. It's to make sure that there's like each new input into a hash function will totally scramble what happens. For neural networks, of course, they need to consider how this piece of information changes what you should make of this other piece of information.

2:04:38And that's an extremely interesting point. I guess at a high level, the difference in what they're trying to do, in some sense, they're trying to do the inverse thing, which is cryptographic protocols are trying to take information which has structure and make it look indistinguishable from randomness. Yeah. And neural networks are trying to take things which are, look like random, protein sequences, DNA, garbled text, and extract higher level structure from it. So they have similar high level mechanisms, but they're actually kind of trying to do the opposite things. Yeah, I wonder what you're making of that.

2:05:15Yeah. So, I mean, like the mixing, like 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 you're making a cake and you want to 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 um but beyond that like in back to the digital world um there are some differences and the one you talk uh call out is is a pretty strong difference um the way it shows up um like what makes neural nets, like if you just randomly initialize a neural network, actually maybe it's a reasonable cryptography, like a cipher as well, because like the random initialization is going to jumble stuff in a complicated way.

2:06:02It 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 like not overcomplicate the derivative. So the residual connection keeps it contained and simple. And so does 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, in theory, in the field of real numbers.

2:06:46and 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's quite difficult to make the difference of the output be small like the whole job of a well-designed cipher is to make the difference in output very large. So I guess the distinction is that the optimization goals at that point are about complexifying. They don't have the same residual connections or like layer norms. Yeah. I mean, I guess a place where the two merge is backdoors.

2:07:31Okay, so with a backdoor 011, you're trying to hide... What do you consider an input? It's not an input into the forward pass, but it's an input into the backward pass, but you're trying to hide an input into the backward pass. Like this is like an adversarial. Yeah. So yeah, I mean, in fact, 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? 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.

2:08:09But that is the common case in ciphers, whereas that's the undesired case in neural nets, for sure. Okay, so I was asking you, have neural networks actually been used for cryptography? And you realize it might be better if you just do this on the Blackboard? Yeah. So I'm curious. Are they actually being used for cryptography? Yeah, so using neural nets for cryptography, 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 were broken. So probably a bad place to start. But the other direction has been very, like in at least one very clear case, quite productive.

2:08:50So there's this construction in, so a construction that exists in ciphers and then was imported into neural nets called a Feistl cipher, Feistl network. So the idea is that you may have some function f, which is not invertible.

2:09:09But you like the function because it does interesting things. Like it does an MLP, for example. Or it mixes it in an interesting way. You'd like to build something out of this that is invertible. So the construction we're going to make is going to actually be a two-input function rather than a one-input function.

2:09:28And we're going to apply ffx we need to actually remember what x was so we're going to stick x over here so that we can work backwards and then we also can't drop y so we're going to remember y and we're going to add them together and so we form this topple so um 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 recover y by z minus f of x, because I've already recovered x. So that means that this construction is invertible.

2:10:15This was used in ciphers, like, a ton. Still is used. It's one of the main mechanisms of constructing ciphers. Often you want ciphers to be invertible, especially the layers of ciphers you want to be invertible, because that has better cryptographic properties. This is actually being ported over into neural nets. There's a 2017-18 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 invariable with exactly this construction.

2:10:56And 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. 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.

2:11:38And 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 0, 1, 2, 3 order.

2:12:22I have to write all of the activations to HBM. And so I get an HBM footprint here that is kind of like linear in a number of layers. Yep. So this actually can be the largest memory footprint during training. And so this is normal training. And then I run the backwards pass and I read it kind of in reverse. Like I run them sort of forward pass goes forward, backward pass goes backwards, and I have to read them back out. The idea of this Revenet's 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 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 have the activations that I need here.

2:13:17So this ends up being a memory saving, which is a nice idea. Interesting. And in some sense, you're spending more compute to save memory. That's right. Yeah. Interesting. Huh. Actually, it's kind of the opposite of what you're doing with the KVCache. The KVCache. Yeah. Yeah. You're spending more memory to save compute. Yeah. Spending more memory to save compute is generally profitable, given where hardware is at. Yeah. Interesting. Cool. That was super fun. Right. Yeah. Thank you so much for doing it. I feel like it really vindicated the vision behind the studio and the Blackboard. Cool, thanks so much for doing it.

2:13:49Thanks.