Recurrence and Attention for Long-Context Transformers with Jacob Buckman - #750

26:55But state size is just as important. It's just much less understood and much less studied because not many people have yet adopted this perspective of viewing the K-V cache of a transformer as its state. But once you do, you see, yeah, like five orders of magnitude of scale between a transformer and an LSTM make it completely impossible that the LSTM is going to compare in performance. It's just not going to, right? It's so far away. So actually, all these interventions that are about reducing the state of a transformer, none of them apply for RNNs, including state-based models, Mamba, all those good architectures, because their states are way too small already.

27:33You want to make those states bigger. And that's actually one of the main focuses of our research at ManFest. So you talked a little bit about the importance of state relative to the number of parameters, which is kind of a more popular way of thinking about model performance. is it just state that we should be thinking about independently of other things that are happening in the model? Or is there more that is required to really characterize state as a relative metric for comparing models? All axes of scale are important. I don't want to say that state is the only important one at all. What you really want is an architecture that's balanced, right?

28:19Just like I was saying before. In this field, it's all about balance. That's how you know. When a model is beautiful and balanced and elegant, that's how you know it's going to be a good compute optimal model. And the problem with having a state that's really big or really small, fundamentally, is that it's imbalanced. So a way to think about this is as the ratio between the flops spent on parameter-based calculations and on state-based calculations. So to paint a picture a little bit here, right? If you're doing, say, the forward pass of a transformer, What you're doing is you're starting with some tokens, you embed them into activations, and then those activations get transformed and transformed and transformed again until they're finally giving you some output.

29:01And each step of transformation costs you something computationally, some number of flops. And so what we can do is we can partition the type of flops into sort of two broad categories that together account for like 99.99 % of the total flops done in the forward pass. Some of them we can say are wait flops. A wait flop is just a multiplication between a parameter and an activation. So in the MLPs, all you're doing are wait flops. Also in the projection between, you know, like when you're projecting X to Q, that's going to be wait flops there. But then there's the state flops. State flops are multiplications between some state element, which for the case of a transformer would be some member of the KV cache and an activation.

29:43So when you're doing, for example, the QK inner product, this is going to be something that is all state flops. And it's the ratio between these two things, how many state flops are we doing and how many weight flops are we doing, that really guides you to compute optimal architectures. And the reason is this, right? Both parameters and state scale log linearly. loop. That means that doubling them and doubling them and doubling them always has the same linear effect on the ability of the resulting model. But of course, the cost in flops is not log linear. The cost in flops is just linear. So doubling either parameters or state will cost you twice as many of that sort of flop and produce a linear improvement.

30:29And so what this means is, if you think about an architecture that's wildly imbalanced, let's say you're using like an LSTM, which has just this tiny little state and a huge amount of weights, relatively speaking. Now, here's two things that you could do that would improve the performance the same amount. You could either double the state size or you could double the weight size. Since they're both log linear, these will have, you know, I'm obviously being a little hand-wavy here, but they'll have roughly the same effect on the outcome. But they will not have the same effect on the computational cost.

30:59Doubling the state size is negligible. Your weights are 100 times larger than your state. So doubling the state, you don't even notice the difference in runtime, but you do notice the difference in performance. So essentially it's a free win, right? If one of these things is much larger than the other, double the smaller thing, free win. And then you apply that same argument again. You can still double it and double it and double it and double it. And eventually you'll get into a regime where you notice the doubling in terms of all cloud time, but that won't happen until weights and states are doing roughly the same number of flops, not exactly one-to-one, but at least within an order of magnitude.

31:31And the same argument applies, by the way, in the reverse direction. So if you have a transformer with a very long context, attention is so, so expensive. You're doing all those state flops. That is basically free to double the weights. You might as well do it, right? And so this means that the only architectures for which it makes sense to train a transformer on a long context have this unfathomably huge number of weights beyond even the very largest models that people are commonly training today to get to a context length of a million or something, right? You need several orders of magnitude of increase in parameter skip, which I mean, maybe we're getting there, but we're certainly not there now.

32:06Does this imply that you could like graph all of these models in terms of the state weight flop ratio and like find some Pareto frontier and it's consistent with what we're seeing in terms of leaderboard performance or benchmark performance, I guess? So I actually haven't done this for sort of like the models that are out there, mostly because it's very hard to get information on exactly what the architectures look like. like and to compute these things. But I have run sort of like smaller, more controlled experiments just on our own experimental setup and seen exactly what you're describing. Essentially, I can show you an experiment in our paper.

32:43I can link it after if you'd like where we trained three models with the exact same everything. They have the same flops per update. They have the same tokens per update, the same batch size, just exactly the same in all regards. The only way they differ is in the ratio, the weight state flop ratio. We have one model that has a huge, huge context and a very small number of weights. Another model that has a huge number of weights and a very small context. And finally, the Goldilocks model, right? In between, perfectly balanced. And you see that, yeah, that's the one that has the best performance.

33:17That's the one that you want to be training, given this sort of like compute budget and these set of constraints. So this is a pretty ubiquitous principle. You'll basically see it everywhere. And yeah, I think if we were able to get details on for each model, how much it had been, how much state it has versus how much weight it has, we'd see this as well. Now, unfortunately, like those aren't the only factors. So you could have two models, one of which is much more well balanced than the other. But maybe the worst balanced one has been trained 100 times as long or on much better data, right? And these will completely mess up the result here.

33:53And so you'll see something very different in downstream performance. That's why I prefer to do it on this nice, clean experimental set. But I think the principle is very robust, and you'll see it everywhere when performing appropriately controlled experiments. So in publishing this power retention architecture, there's the paper and kind of a mathematical description of it. There's a model that you published. which is a 3 billion parameter coding model, if I'm remembering correctly. And the third piece is some CUDA kernels. Can you talk a little bit about how these elements, you know, fit together and maybe your general experience of, you know, building this out beyond the, you know, the broader perspective that you, you know, gained into the problem?

34:48So we are a very practical research lab. We aren't satisfied with just getting a theoretical win. We want to actually have the best architecture, the one that is most worthwhile for almost everyone to train. Basically, anyone who's doing long context, anything should be training their models using our architecture. That's our goal. And the first piece of that is, as you said, mathematical. We need an architecture that actually, when measured per flop, gets the performance gains that we expect, that it actually is able to do fewer operations to produce more intelligence. And that's really what the paper outlines.

35:23It shows very clearly, yes, when you focus on getting models that have a balanced weight state flop ratio, and you use sort of models that process the entire input instead of something like sparse attention that's only processing a piece of the input. That's like this. Those are the two things that are required to get the best performance, the best scaling loss. Process the whole input and do so with a balanced weight state flop ratio. And that's the first piece that we explain in the paper with very careful comparisons to transformers, ablations, really just painting this whole picture that we've been talking about this whole time.

35:59But that's just theoretical. Next piece is it's gotta be practical. We need something that you can actually run and have wall clock speedups. And this actually, there's a little bit of an aspect of co-design between this part of things and the mathematical thing I mentioned a moment ago, because we would only be happy with an architecture if we knew that it could be given a very efficient GPU-friendly implementation. But the second piece is to actually make that GPU-friendly implementation. We started off with Triton, which works fairly well. You're able to get very good utilization. The kernels are written basically in a similar style to Flash Attention in Triton, but that's not the most you can do.

36:40So we actually went to level deeper. We knew that we could really, really crush Flash Attention if we squeezed all the performance possible out of the hardware. And because this architecture, some of the pieces of it are a little unique, a little complex to do, we weren't able to do it at the Triton level. We had to drop down a level of abstraction and just program directly in CUDA. And it was so complicated that even typical CUDA programming patterns didn't really work for our application. They just, there were too many moving pieces, too many like sort of like workarounds that we needed to do.

37:11So we realized that actually the best way to write these kernels is to build a whole CUDA framework from the ground. So we built something that we call Vidrial, which is basically a way of writing non-spaghetti CUDA, not to put too fine a point on it, right? Everyone who writes CUDA kernels does so. There's very little code reuse, no testable components, a bunch of if statements everywhere with different compilation paths, hundreds and hundreds of lines of very complicated, nuanced pointer arithmetic in order to make all the memory movements happen correctly. And if you ever want to make a change to maybe like exploit new hardware or adapt to a different problem shape, it's a nightmare.

37:47Like is so much rewriting and retesting in order to make that happen. So Vidrial is a framework that gives you just a set of very clean patterns for writing beautiful, testable, and very, very efficient CUDA. Basically what Vidrial does, it lets you write every possible CUDA kernel that implements this operation. So always, all these are going to have basically, they're going to be doing the same thing, but they're going to be using the hardware in very different ways, in all the different possible patterns that one might consider using. Things like tile shapes, right? How much memory does each thread own and copy?

38:22How much memory, like where does it copy it to? Like which thread does it hand it off to? Things like pipelining. How are operations being staggered within the GPU? All of these things can have lots of different values and you have to basically just do a huge empirical sweep to find out which ones are the best for any given setup. Visual lets you just write the generic patterns that you want, and then gives a clean Python JIT system for doing any empirical sweep. And to be clear, is it specific to power architecture implementation, or is it a generalized framework for building CUDA kernels? It's completely generic.

38:58So we built it because we wanted to write kernels for power retention, but in order to write those kernels, we realized we just needed to make a good general CUDA framework, which we did. And we implemented lots of other kernels as well, things like basic addition, like element-wise operation, matrix multiplication, copy kernels. We also implemented flash attention just to see what would happen. And we actually get 20 % speedups over tree DAO's flash attention on a whole bunch of problem shapes, which we thought was pretty cool, pretty good proof point for the framework. And of course, on power retention itself, we're now able to get 2 to 4x speedups over the Triton implementation because we're really squeezing what can be squeezed out of the hardware.

39:38Yeah, it's interesting. You could easily see that going the other way, adding layers of abstraction and cleanups and organization, having a computational cost. But it sounds like the typical CUDA kernel implementation is fairly messy and has some inefficiencies that you're able to clean up a bit systematically. Exactly. And I'm very, very proud of what the team is able to put together here because it is exactly what you said, right? You usually think about going up a level of abstraction as introducing inefficiencies. But in this case, it actually allows us to search the space of possible low-level implementations to get the best one.

40:16And that's something that can't be done without good abstractions basically partitioning what the different useful low-level implementations are. We rely heavily on something called QtLayout, which is a sort of like recent addition to the Cutlass library by the team at NVIDIA. That's absolutely brilliant innovation. And I think we're probably the most advanced users of Qt out there because our entire system for basically computing which memory needs to be moved where across basically any different possible configuration you can imagine relies on these sort of layout manipulation operations in Qt.

40:54Interesting. Have you been tracking what Chris Latner and Modular are doing with Mojo? they I think raised some money this week and I think of Mojo as a well I initially thought of it as another language but it's kind of like bringing more CUDA flexibility directly into Python like semantics I guess that's the way I think about it have you been following that? I think the work is amazing we haven't gone deep on Mojo yet I think we wanted to stick with CUDA because it's the most mature and we just wanted some really fast-working kernels instead of trying to explore other frameworks. But we really like the idea of Mojo and hope to, at some point in the, hopefully near future, generalize the Vigil framework to also work with Mojo as a backup.

41:46Ultimately, the core principles that we're leveraging are completely agnostic to what's going on underneath. It's done in CUDA and in C++ via C++ template metaprogramming because that's the current best practice way to do it. But the core ideas, really the two core ideas, one is this separation of static and dynamic computation where you at static time, like at compile time, learn everything you can about what memory is moving where for any given configuration. And when you say learn, do you mean learn as a human developer or learn as an optimizer a software optimizer that is um you know building out the low-level components yeah and learn in the software sense the it's not even like a machine learning maybe that was a bad choice of words honestly um really it's just uh you just you just infer it just at static time profiling kind of similar to this yeah you can say, okay, like if the problem shape is this and the layout is this, and we're going to want to use a tile size of that and break it up into this many threads, then okay, that means this thread needs to take this set of memory addresses and move them to here.

43:01And this can be known statically, right? It doesn't matter what data you're giving the kernel. You just need to know the shape that you want in, the shape that you want out, and what we call the configuration of the kernel, which is basically the way that it's getting executed on the underlying hardware. And with that information, you can just learn a whole bunch about how this computation is going to go. And then when you give it some data, it just executes all those operations and you're done. So that's like one of the core ideas. And the other is this integration with Python. So we want it to be very easily usable, very easily suble.

43:34We have this JIT system. So you can just in time compile on the fly, a brand new, hyper-specialized visual kernel for any given problem shape, any given hardware platform. And by just running a sweep over different possible configs, you can just empirically measure each one and find the best for your very specific problem shape. So our gains over flash attention on normal problem shapes, the ones that the authors of flash attention expected people to be using, our gains there are minimal. They're in the like one to 2 % range. But where we get the huge speedups is on the problem shapes that are less typical.

44:13Like if you want to have some weird number of tokens in your sequence length. For example, if you're doing pre-fill on a real document, right? That document isn't going to have a nice even 1 ,024 tokens length. It's going to have some crazy number, right? And in that case, our ability to sweep over different ones will give the ideal way to basically execute that operation on the hardware. So I guess this doesn't really apply if you're just doing one document of that length. But if you're doing a whole bunch of documents, all of that weird length, then our ability to like sweep over different configurations, find the best one and run every single one with that best configuration really pays off.

44:50You also, as I mentioned, put out a model. Talk a little bit about the model itself. So this is sort of the third pillar. The first one was it needs to be mathematically and theoretically optimal. The next is it needs to actually run fast on hardware. And the final one is you need to want to train this thing. And in the context of the modern era of transformers, right, every good training run starts from some great initial weights. Unless you're one of the big labs, you're very rarely going to want to initialize your model randomly, right? You're always going to want to start off from some pre-trained weights that already do something useful.

45:26Now, by switching the architecture from using exponential attention to using power retention, this is going to completely mess up the outputs, right? The existing weights are going to produce basically garbage. But by doing a little bit of additional retraining, you can get right back to the original performance. It's way, way cheaper to start from the bad predictions of a good set of weights on a power retention architecture and train it to match the performance that those weights were able to get on exponential architecture than it is to train from scratch. So in other words, StarCoder 2, This is a 3 billion parameter model trained by the big code project to be a fast local coding assistant.

46:11Has good performance. I think, for example, it has like 30 % accuracy on human eval, which is a Python coding eval. If you just take that, take those same weights. So this is a tune that you did for the power retention architecture that you are offering as a starting place as opposed to a ground up model that's based on power. You don't need to build a ground up model to use power retention. You can tune a model that you like. Exactly. So yeah, Power Coder is this tune of Star Coder. And it sort of has two things. One is it's just a useful model. If you want a model that is a coding assistant at the 3 billion parameter scale and you want it to have really fast inference, use PowerCoder.

46:55It's just better than anything else out there. Now, the other thing it is is a proof of concept, right? It just shows that, hey, if there is a model that you care about and it's pre-trained as an exponential transformer, all you have to do is do a little bit of retraining. For example, we were able to recover performance from StarCoder 3B to PowerCoder 3B. The sort of like realignment period, we call the metamorphosis, took around two hours on a cluster of 128 H100s. It's not nothing, but it's not massive. It's not something that you need big lab pre-training scale in order to pull off. And is there a complex recipe required or is it fairly straightforward tune?

47:35Surprisingly, no complex recipe is required. We did do various explorations with the recipe to see if we could squeeze a little bit more out, and you can. But the basic recipe is like 99 % of the way there. Literally just set it up like a training run. A training run like any other with the power retention architecture, but instead of initializing randomly, initialize to the known good star coder weights. And after two hours of training, you've recovered 30 % score on human eval, and with more training, you can push it even beyond. So it's actually just a really, really easy way to turn any existing model into the power retention bearing.

48:12Okay, I'm not sure I understand it. The known good star coder weights, where do those come from? Like how's that different from star coder? So if you just download the star coder weights and the star coder architecture, you can run star coder. If you download the star coder weights and download the star coder architecture, but then replace the call to flash attention in the star coder architecture with a call to power retention, You now have a new model with weights that are equivalent to the StarCoder weights. And if you take this model and train it for just two hours, it'll recover the performance that StarCoder initially had.

48:44So what I haven't heard you talk about in terms of performance, at least explicitly, is long contexts. Like, is the testing that you've done to demonstrate with, you know, existing benchmarks focused on, you know, small contexts or a normal context, like you're performing better with Starcoder? or can you take this star code of 3b parameter model and demonstrate that you can then, you know, shove a much longer context in it? I don't think that works because the model architecture kind of fixes the context, right? Yeah, but you can just change. So there are some parameters like the rope, the rotational positional embedding parameters that do sort of fix the context length, but you can work around it.

49:35You can just tweak those two and do additional retraining and that also works fairly well. And so, yeah, this is basically exactly what we do. We start, like our initial training run is usually to recover performance. So StarCoder 2, for example, trained at a context length 4K. So our initial performance recovery just happens at context length 4K on the same data that StarCoder was trained on. But then, in fact, this is currently an ongoing run. We'll be releasing updates to the way it's due when it's completed. You can actually just train it on 16K and do another big training run on 16K. And now, because you're just paying a linear cost per token of context instead of a quadratic cost, you can actually have that be like a tractable, almost like a cheap training run.

50:16So way cheaper than training StarCoder on 16K, for example, is training PowerCoder on 16K. You could push it to 32K, 64K. At this point, it just becomes a resource question, right? How much do you want to invest in these long context-free trainings? and we're a small lab. So, you know, we make do with what we have and are pushing it as long as is reasonable given our budgets. But we're really excited to see as this sort of starts to pick up momentum and other people get excited about it too, what models other people train and how far they're able to push the context. Presumably you feel comfortable that it works without having demonstrated the ability to push the context, you know, out very long.

50:56Like what gives you that confidence? I guess that's where the mathematics come in. But, you know, what do you talk a little bit, talk a little bit more about like the practical realities of, you know, pushing context out? Like how, you know, what do you think is possible? Maybe that's a way to ask the question. What do you think is possible with this model? And where do you think there will be limitations? I think the optimal context length is always going to be a scaling law, just like every other axis of scaling. And it's one that's going to be highly dependent on the data set. So I think that this is something that a lot of the big labs are doing these days, you know, just putting these big numbers on this is my max context length, while not actually verifying that the performance within that context length is any good, right?

51:50And of course, as we were saying at the beginning of the chat, that's actually the most important thing, is that it's actually able to use those tokens of context. So by looking at these in-context learning curves, where you basically are able to see how much it's able to use each additional token of context, you can really deeply learn something about whether this model is taking advantage of the extra information it was given. And what you find is that on any given computational budget, there's a particular context length that is best suited for that computational budget, just like there's always a particular model size that's best suited for it.

52:25And it's these scaling laws that really give us confidence. Scaling laws themselves are so well established at this point on all these other domains, on model size, on batch size, on training steps, right? There's no question that scaling laws work. Now, state size scaling laws, which basically govern the performance as you increase the state, and context length scaling laws, which same as you increase the context length, these are less well understood. And I think a big part of this is because in transformer models, these two things are tied together. Increasing the context length means increasing the state size.

53:00It's the only way to do it and vice versa, right? So nobody has really, before us, carefully isolated these two factors and explored their independent effect. But they do have independent effects, right? And you can really reliably see these scaling behaviors. When you look at the plots, they look just like all the other scaling plots that we know and love and trust deeply. So I think really it's this more than anything else that gives me confidence that this is sort of a fundamental general principle and one that's going to just like very naturally be able to achieve and maintain good performance as the context length gets longer.

53:34I think there's a lot of other approaches, for example, windowed attention, where you're just like arbitrarily truncating the attention at a certain size that we've measured. And we've looked at what the in-context learning curves look like for windowed attention. And guess what? You see that it learns quickly at the beginning. It's able to use the first couple tokens of context, and then it plateaus. right and that long plateau where it stops improving as you show up more tokens is exactly what you would expect for a model that is only looking at a small recent window of tokens and interestingly the the point at which this knee happens and it stops learning more is way shorter than what people call the effective context length of windowed transformers which is like the depth times the window size so really window transformers are just not able to use very much of their context.

54:23But ARM architecture is. And we think that as you continue to scale it up, all these trends will just continue and we'll get the behavior we expect, even in very, very large state sizes and very, very long contexts. And of course, very large parameter counts as well. Of course, as a small lab, all we can do is extrapolate, right? And we're currently working on progressively scaling up. And every time we do, we see that things match what we expected to see, and then we move on. But ultimately, a big part of the reason why we're open sourcing this work and excited to share with the community is so that lots of people can play with it themselves.

54:57You know, we really deeply believe in this research. Based on the results that we've seen, we're fully convinced that this is something that actually works. And so we're not afraid to put it into everyone's hands and let everyone see for themselves that it actually works. And we expect this isn't something that we're going to do alone. This is something that we're going to, you know, basically do in collaboration with the broader community. For folks that are intrigued and want to play around with it, where do you expect them to start or where do you suggest that they start? And what do you want them to know or think about as they're getting started?

55:30I think a great place to start would be with the PowerCoder model. It's available on Hugging Face. You can just download it and use it. You can play around with the inference. You can see the really fast inference speeds. You can extend the context and see the inference speed is consistent per token. I also think that it would be fun if people want to do some additional post-training. For example, this is not an instruction-tuned model. It's just trained on raw code, but maybe a little instruction fine-tuning would be a cool project just so you can see this thing actually learns just like a transformer would in similar situations.

56:03Beyond there, I would love it if people who are already using open-source models, especially people who are already fine-tuning or post-training open-source models for their own needs, would consider doing a metamorphosis run. where they take their favorite model, maybe it'll be a QuenCoder, 30 billion QuenCoder or something, and rather than post-training it directly on their preferred post-training data set, they first metamorphose it into power retention and then do the post-training. And if their preferred post-training data set is something that has very long context, then they'll see big, big speedups.

56:37We'll get like, for example, 64K context, you'll get a 10X speedup at train time and a 100X speedup at inference time. just by switching over to power retention. And yeah, I think just comparing those learning curves will let people convince themselves that this is something that is actually really going to be useful for them. Awesome, awesome. Well, Jacob, thanks for jumping on and sharing a little bit about what you've been working on. Super interesting stuff. Thanks so much for having me. Very fun conversation. Awesome, thank you.

57:14Thank you.