FlashAttention-2: Making Transformers 800% faster AND exact

Today we have Nuswix, um, because he's in, he's in Singapore, so, uh, it's a, it's a one-on-one discussion with, uh, Tri Dao. Welcome.

Hi, everyone. I'm, I'm Tri Dao. Excited to be here.

So Tri just completed his PhD at Stanford a month ago. Um, you might not remember his name, but he's one of the main authors in the FlashAttention paper, which is one of the seminal work in, uh, the transformers era.

Um, he's got a lot of interest from efficient transformer training and inference, uh, long range sequence model, um, a lot of interesting stuff. And now you're gonna be a assistant professor in CS at Princeton next year.

Yeah, that's right. Yeah.

Nice. And in the meantime, just to get, you know, a low pressure thing, you're a chief scientist at, at Together as well, which, uh, is the company behind the Red Pajama.

Yeah. Yeah. So I just, uh, joined, um, th-this week actually, and it's been really exciting. Yeah.

Nice. Uh, so what's something that is not on the internet that people should know about you?

Hmm. Let's see. Um, I think, uh, before c- when I started college, I thought I was gonna be an economist.

Hmm.

So I was, uh, fully on board. I was gonna major in economics. But the first week I was at, at Stanford undergrad, I took a few math classes, and I immediately decided that I was gonna be a math major, and that kind of, um, changed the course of my career.

Wow.

So now I'm doing kind of math, computer science, AI research.

Nice. That's, uh... You know, I had, I had a similar thing. I started with, uh, with physics, and then I took like a programming course-

Mm-hmm

... and I was like, "I gotta do computer science." "I don't wanna do physics." Um, so FlashAttention is definitely, you know, everybody's using this. Everybody loves it. You just released FlashAttention-2-

Yeah, that's right

... last week. Um-

Yeah, early this week, on Monday. Yeah.

Yeah. And, you know-

AI time.

Things-

Yeah, things move fast.

Yeah. Four days ago is one week ago- ... in, in AI. Uh, so maybe let's run through some of the FlashAttention highlights.

Mm-hmm.

Um, some of the innovation there.

Yeah. For sure.

And then we can dive into FlashAttention-2.

Yeah.

So the core improvement in FlashAttention is that traditional attention is a quadratic-

Mm-hmm

... sequence length.

Yeah.

So it's n to the, to the two. FlashAttention is linear, which obviously helps with, uh, with scaling some of these models. Um-

Right. So, so there, there are two factors, um, there. So of course the goal has been to make attention, um, go faster-

Mm-hmm

... or more memory efficient. Um, and ever since attention became popular in, um, 2017, uh, with the transformer paper, lots and lots of folks have been working on, on, on this. Um, and a lot of approaches has been focusing on approximating attention.

Mm-hmm.

Um, the goal is you wanna scale to longer sequences. Um, there are tons of applications where you wanna do that. Um, but scaling to longer sequences is, is difficult 'cause attention-

Mm-hmm

... scale quadratically in sequence length and both runtime and memory, as you, as you mentioned. Um, so, um, instead of trying to approximate attention, we, we were trying to figure out, you know, can we do the same computation, um, and, uh, maybe be more memory efficient.

So in the end, we ended up being the memory is linear in sequence length.

Mm-hmm.

Um, in terms of computation, it's still quadratic, but we managed to make it much more hardware friendly, and as a result, we do get, um, wall clock speed up on the order-

Mm-hmm

... of two to four X. Um, which really helps 'cause that, that just means that you'll be able to train with two to four X longer sequence length for the same cost, um, without doing any approximations.

Right.

So as a result, um, lots of folks have been using this. Um, thing is available in a lot of, um, libraries that train, um, that train do, um, language model, um, training or fine-tuning.

Mm-hmm. Yeah. And the approximation thing is important because this is a exact thing versus-

Yes

... like a sparse.

Yeah.

Um, so maybe explain a little bit the difference there.

For sure. For sure. Yeah. So, um, you know, u-attention, um, essentially you compute pairwise similarity between, um, every single element in a sequence-

Mm-hmm

... uh, against each other. Um, so there's been other approaches where, um, instead of doing all that kind of pairwise computation, you only compute similarity for, you know, some pairs of, of, of, um, of elements in the sequence.

So you don't do kind of quadratic number of comparison. Um, and, and this can be seen as a k- as some form of sparsity.

Mm-hmm.

Es-essentially, you're, you're ignoring some of the elements. When you write down the matrix, um, you essentially say, "Okay, I'm gonna pretend they're zero."

Yep.

Um, and, and that could, uh, that c- uh, that has some benefits in terms of runtime and, and memory. Um, but the trade-off is that it tends to do worse in terms of quality-

Mm-hmm

... um, 'cause you're essentially approximating or ignoring some, some elements. Um, and I personally have worked on, on this as well for, uh, a few years. But when we talk to, uh, practitioners who actually train models, especially at large scale, they say, "Well, we, uh, tend not to use these, um, approximation-"

Mm-hmm

... "approximate attention methods." Um, because turns out, this, this was surprising to me, um, at the time, was that, um, these approximation methods, uh, even though they perform fewer computation, they tend to not be faster-

Hmm

... in wall clock time. So this is-- it was pretty surprising because back then I was, um, I think my background was more, uh, on the theoretical side.

Uh-huh.

So I was thinking of, oh, how many FLOPS or floating point operations, um, are you performing? And, and hopefully that correlates well with wall clock time. Um, but I realized that I was missing a bunch of ideas from the system side where, since, uh, uh, FLOPS or floating point operations don't necessarily correlate with runtime.

There are other factors like memory reading and writing-

Mm-hmm

... parallelism, and so on. So, um, I learned a ton from just talking to systems people because they, they kinda figured this, this stuff out, you know, a while ago. So that was really, uh, eye-opening, and then we, we ended up focusing a lot more on memory reading and writing-

Mm-hmm

... because that turned out to be the, um, the majority of time when you're doing attention is, uh, reading and writing memory.

Yeah. Yeah, the I/O awareness-

Yes

... is probably like one of the biggest-

Mm-hmm

... innovation here. Um, and the idea behind it is, like you mentioned, the, the FLOPS growth of the cars have been going up, but like the, the memory bandwidth not as much, so, um-

Yes

... I think like maybe that was one of the assumptions that the original attention paper had.

Yeah.

Um, so talk a bit about how that came to be as an idea, you know? Like how-

Yeah

... it, you know, it's one of those things that like in insight, it's like obviously why are we like rewriting to like-

Mm-hmm

... HBM every time.

Yes.

You know?

Yeah.

And like once you, once you change it, it's clear. But-

Mm-hmm

... um, what was that discovery process?

Yeah. Yeah. So I think in, in hindsight, um, a lot of the ideas have already been there in the literature. Um, and I would say is, um, it was somehow at the intersection of both, uh, machine learning and systems, um, and you kinda needed ideas from both, both sides.

Mm-hmm.

Um, so on one hand, um, on the system side, so lots of systems folks have, have known that, oh, you, you know, kernel fusion is, is great. Uh, kernel fusion just means that instead of performing, you know, loading, um, the same element, um, and, uh, instead of performing, um, uh, an operation, write it down, load it back up, and perform the second operation, you just load it once, perform two operations, and then write it down again.

So that saves you-

Mm-hmm

... kind of, uh, memory read and write in the middle there. So, uh, kernel fusion has been, um, uh, a classic. Um, there's been other techniques from the system side like, um, tiling, where you perform things in, in, in, uh, perform computations in block, again, so that you can load it into a really fast memory.

Think of it as a cache. Um, and this is again, classical computer science ideas, right?

Mm-hmm.

You wanna use the cache. Um, so the system folks have been thinking about, um, these ideas for, for a long time, and they've, they apply to, uh, attention as well. But there were certain things in attention that made it difficult to do in a complete kernel fusion.

Um, one of which is there is this softmax operation in the middle, which requires you to essentially sum across the, the row of the attention matrix.

Mm-hmm.

So it, it makes it difficult to kind of break it, because there's this dependency, so it makes it difficult to break things into a block. So on the system side, it's been, um, people have been thinking about these ideas, but it's been difficult to kind of do kernel fusion for the entire-

Mm-hmm

... operation. On the machine learning side, people have been thinking more algorithmically. They say, "Okay, either we can change the, uh, we, we can approximate attention," or, um, there, there's this trick called the online softmax trick, which says that you can-- Because of softmax, um, the way it's written mathematically, you can actually break it up into smaller pieces-

Mm-hmm

... do some rescaling, and still get the right answer. So this online softmax tri-trick has been around for a while. I think there was an paper from NVIDIA folks back in 2018 about this, and then there was a paper from, um, Google.

So, um, Marcus, Rob, and, and, and, and, and Stutz wrote a paper late 2021 on using this online softmax trick to break attention, uh, up into smaller pieces.

Mm-hmm.

So a lot of the ideas were already there. Um, but turns out, um, I think if you, um, you kinda need to combine ideas from both sides. So you need to understand that, hey, we wanna do kernel fusion to reduce memory reads and writes, but we also need this online softmax trick to be able to-

Mm-hmm

... break the softmax into smaller pieces so that a lot of the systems tricks, um, kinda carry through.

Right.

And, and so, um, you know, we, we, we saw that and, and it was kind of a natural idea that we, we, um, we ended up using ideas from both sides, and it ended up, uh, working pretty well.

Yeah. Are there any downsides to kernel fusion? Uh, if I think about databases and the reasons why we have, um, atomic operations, you know, it's like you have observability and fallback in between them. Um, yeah, how, how does that work with, with attention?

Like, is there anything that we lose by, by fusing the operations?

Yeah. I think mostly on the, kinda on the practical side is that, um, when you do kernel fusion, it's a little bit, uh, you lose a little bit of flexibility in the sense that, hey, now you have, um, for example, it's, uh, FlashAttention is just a subroutine that you would call to do attention.

Um, but as a researcher, let's say, um, you don't want that exact thing, right?

Mm-hmm.

You don't want just attention. Let's say you want some modification to attention. You wanna do, "Hey, I'm gonna multiply the query and key, but then I'm gonna do this extra thing before I, you know, carry on." Um, and so kernel fusion just means that, okay, in, we have a subroutine that does the entire thing.

But if you wanna ex-experiment with, uh, with things, you, you won't be able to use that, um, that fused kernel.

Mm-hmm.

And, um, of course the, the, um, the answer is, can we have a compiler that then automatically does a lot of this, uh, this kernel fusion? Um, and, um, you know, lots of compiler folks are, are, are thinking about this, um, either with a, a new language or with, um, you, you, you, you can embed it in PyTorch.

So the PyTorch folks have been working on this as well. So if you write just your code in, in, in PyTorch and the, and, and they can capture the graph, can they generate, um, code that will-

Mm

... kind of fuse everything together? Now, that's still un-ongoing, and they-- It works for some cases, but for attention, because of this kinda softmax rewriting, um, stuff, it's been a little bit more difficult. So, um, maybe in a year or two, we'll, we'll have compilers that are able to-

Mm

... do a lot of these, uh, optimizations for you, and you don't have to, for example, spend a couple months writing CUDA to get- ... this stuff to work.

Awesome. And just to make it clear for listeners, when we say we're not writing it to memory, m-we are storing it, but just in a faster memory.

Yeah.

So instead of the HBM, we're putting it in the, in the SRAM.

Yeah, that's right.

Um-

Yeah

... yeah, maybe explain just a little bit the, the difference there.

Yeah, for sure. So, um, this is, um Uh, you know, kind of a caricature of how you, you think about, um, accelerators or, or GPUs in particular is that they have a large pool of memory, usually called, um, HBM, a high bandwidth memory.

So this is what you think of as GPU memory.

Mm-hmm.

So, you know, you're using, um, A100 and you, you list the GPU mem- memory is like 40 gigs or 80 gigs. So that's, that's the, that's the HBM. Um, and then, um, when you perform any operation, um, you need to move data from the HBM to the compute unit, so the, the actual hardware unit that does the, the computation.

And next to the-- these compute units, there are, um, uh, it's called on-chip memory or, or SRAM, which, uh, are much, much smaller than HBM, but much faster. So the analogy there is if you're familiar with, say, CPU and RAM and so on.

So you have a large pool of RAM, um, and then you have the CPU performing the computation. But next to the CPU you have, you know, L1 cache and L2-

Mm-hmm

... cache, which are much smaller than, than, than DRAM and much, but much faster. So you can think of SRAM as like small and fast cache that's, uh, stays, uh, close to the compute unit. Like physically this-

Mm-hmm

... is, is closer. Um, and so, uh, there is some kind of asymmetry here. So HBM is much larger, um, and SRAM is much smaller, but much faster. And, um, one way of thinking about it is how do-- how can we design algorithms that take advantage of this, uh, asymmetric memory hierarchy?

Mm-hmm.

And of course, like lots of folks have been thinking about this back, uh, back in the, I think 1980s when people were, um... Yeah, these, these ideas are, were, are, are pretty old. So ba- I think, um, back in the, uh, 1980s the, uh, primary concerns were, uh, sorting.

Mm-hmm.

How can we sort numbers as efficiently as possible? And the motivating example was banks were trying to sort their transactions and that needs to happen overnight so that the next day they can-

Right

... they can be ready. And so, um, the same idea apply, which is that they have slow memory, which is, was, uh, which was, uh, disk, like hard disk, and they have fast memory, which was DRAM.

Mm-hmm.

And people had to design, uh, sorting algorithms that kind of take advantage of this, um, asymmetry. And turns out, you know, these same ideas can apply today, which is different, uh, different kinds of memory.

Yeah.

Yeah.

Yeah, and, um, in your paper you have kind of like the pyramid of memory. And just to give people a- an idea, when, when he says smaller, it's like HBM is like 40 gig and then SRAM is like 20 megabytes.

Yeah.

So it's not like a little smaller.

Mm-hmm.

It's like much smaller. Uh, but the throughput on card is like 1.5 terabytes a second for HBM and like 19 terabytes a second for SRAM, which is a, a lot larger.

Yeah.

How do you think that evolved? So TSMC said they hit the scaling limits for SRAM.

Mm-hmm.

They just cannot grow that, that much more.

Yeah.

Um, HBM keeps growing. HBM3 is gonna be 2X faster than-

Mm-hmm

... HBM2. I think the, the latest NVIDIA thing-

Yeah

... has a HBM3 on it.

Yes.

Um, how do you think about the future of like FlashAttention? Like do you think HBM is gonna get faster enough and like maybe it's not as, uh, as useful to use the SRAM or?

Yeah. Yeah. I think, I think that's, that's right. I think it's, it comes down to physics when-

Mm-hmm

... you know, when you design hardware. It's just, um, literally SRAM stays very close to compute units and so you don't have that much area to essentially put the, um, put the SRAM, to put the transistors. Um, and um, you can't shrink these things, um, too much.

Um, so just physics like in terms of area, um, you don't have that much area for the SRAM. The HBM, um, um, technically is, is, is off chip, so there is, um, some kind of, uh, bus that essentially-

Mm-hmm

... transfer data from HBM to, to the compute unit. So you have more area to essentially put these, uh, um, these, uh, memory units. Um, and so yeah, I think in the future SRAM probably won't get that much, um, larger, um, because you don't have that much area.

HBM will get larger and, and, and faster. And so I think it becomes more important to, um, design algorithms that take advantage of this, um, this memory as- asym- asymmetry. It's the same thing when in, in CPU where, um, the, the cache is really small.

The DRAM is grow- growing larger and larger. Um, DRAM could get to, I don't know, two terabytes, six terabytes-

Mm-hmm

... or something. Um, whereas the cache stay at like, I don't know, 50 megabytes or something like that. Um, and so, um, I think maybe the algorithm design becomes more and more and more important. There, there's still ways to take advantage of, of this.

Um, I think so in the future, I think, um, FlashAttention, you know, right now is, is being used. I don't know if, um, in the next couple years some, you know, some ar- some new architecture will, will come in and, and whatnot, but attention seems to be still important.

Mm-hmm.

Um, for the next couple of years I still, I still expect some of these ideas to be useful, not necessarily like, you know, the exact code that, that's out there, but I think these, these ideas, um, have kind of stood the test of time.

The ideas like IO awareness from back in the, uh, 1980s, ideas like kernel fusions, tiling. These are classical ideas that have stood the, the test of time. And so I think in, in the future these ideas will become more and more important, um, as we, as we, uh, scale models to, to be larger, as we have more kinds of devices, um, where performance and efficiency become much, much more important.

Mm-hmm. Yeah. Yeah, and we had, uh, Jonathan Frankel on the podcast-

Mm-hmm

... and, uh, if you go to isattentionallyouneed.com he has an outstanding bet and he does believe that attention will be the, the state-of-the-art-

Yeah

... um, architecture still in a few years. Um, did you think FlashAttention would be this popular? Like I'm always curious on the research side, you know, you publish a paper-

Mm-hmm

... and obviously you know it's great work, but sometimes it just kinda falls flat in the industry.

Yeah.

Uh, could you see everybody just like just starting to use this or was that a surprise to you?

Yeah. So I, I think, uh, you know, um, certainly I didn't anticipate the level of, of, of, uh, popularity. Of course, we- we're extremely happy to, to, um, have people using this stuff and giving us feedback and, and so on, and, and help us improve, um, things.

I think when, when we were writing the, the paper, I remember sending an email to, to, to one of my advisors like, uh, "Hey, I'm excited about this paper, um, but I think the most important thing will be the artifact, which is the code."

Mm.

Um, so I, I knew that, like, the code will be valuable. Um, and, um, and so we, we kind of focus a lot on, on the code and make sure that, you know, the code is, is usable and, and as fast as can be.

Of course, the idea, um, the, the paper presents, um, the, the ideas and ex-explain it and have experiments that validates, um, uh, the idea. But, uh, I ca- I, I, I can knew that the artifact or, or the code was the, um, was also pretty important.

Um, and that turned out to be, um, kind of the, the right focus, which is-

Mm-hmm

... you know, we, we put out the paper, we release the code, and continue working on the, on, on the code with my, um, um... So yeah, it was, it is, it's a team effort with my, my co-authors as well.

Yeah.

Yeah. We mentioned Hazy Research a bunch of times on the podcast before.

Mm-hmm. Mm-hmm.

Um, uh, I would love for you to spend five minutes just talking about how, how does the group work? You know-

Sure

... how do people get together? Like, how do you, like, bounce ideas off of each other?

Mm-hmm. Yeah. Yeah. So Hazy Research is a, uh, is a research group at, at, uh, Stanford led by, uh, one of my advisors, Chris Re. Um, and, um, I, I love, I love the people there. It's one of the best experience I, I had.

Like, they've, they've made my PhD, you know, so much more enjoyable. And I think, uh, there are a couple of, um, there are a couple of ways that, um, the group has been, um, has been working pretty well.

So one is I think there's kind of a, a diverse pool of, of, of people who either... You know, some of them focus on algorithms and, and theory, some of them focus on building systems, some of them focus on applications.

And as a result, uh, there is this flow of idea. So as an example, um, we were working, um, or some of, some of us were working on, like, more algorithms and, and, and theory, and then we can, we can talk to the folks building system and say, "Hey, let's try it out and let's, let's put it in the systems-

Mm

... and, and see how it is." Um, and, and there you will get feedback from systems folk. They will say like, "Hey, we implemented this," or like, "We tried this and, you know, this is where it doesn't work."

Mm-hmm.

Something like that. Um, and once we put it in the systems, like, the application folks can, can use, um, the, the algorithm or new methods or, or new models, and we again get great feedback from them because, uh, the application folks, for example, um, you know, some of my good friends, they focus on medical imaging or seizure detection, and that is the problem they care about.

Mm-hmm.

Right? And if your method doesn't work on the task they care about, they will tell you. Um, whereas I think a lot of, um, people in machine learning, they're a little bit more flexible, so they will be like, "Hey, it doesn't work on seizure detection.

Let's try some other task."

Mm-hmm.

Right? Um, but having that direct feedback of like, "Hey, it doesn't work there. Let's figure out why," um, I think that that feedback allows to do better work. Um, and I think, you know, that, that kind of process of exchanging ideas, um, validating it on, on, um, in a real system so that applications folks can try it out and, and give you feedback, I think that that cycle has been, um, very, very useful.

Um, and so that's, that's one, you know, having a diverse group of, of people. Um, the other one is, um, and this is I-- something I really appreciate, um, from advice from Chris, was, um, try to understand the fundamental, right?

And, and he's happy, um, letting me go, go off and read some textbooks-

Mm

... and, and playing with things, because I think a lot of research, uh, ideas come from understanding the old literature and see how it fits with the, um, with the new kind of landscape. Um, and so if you just read new archive papers every day, you know, that's, that's great, but, um, you also need to read textbooks.

And, uh, and, and that's one advice I, I got from Chris, which is understand the fundamentals, and I think that allows to do, um, you know, more impactful work. Yeah.

Yeah. How do you think about, um, academia versus industry? Like, AI machine learning has been an area where up until three, four years ago, most of, like, the cutting-edge work was being done in academia, and now, you know, there's all these big, um, industry research labs.

Um, you're obviously going to, to Princeton, so you're a academia believer.

Mm-hmm.

Uh, how should people think about where to go? Say I'm like, uh, you know, I'm doing my master's, uh, I have to decide between doing a PhD and, like, going into OpenAI Anthropic. How should I decide?

Yeah. So I, I think they kind of play complementary role, um, in, in my opinion. Of course, like, I also, uh, uh, was, uh, considering different paths as, as well. Um, so I think, um, right now, um, scaling matters a lot es- especially when you talk about language models and generative AI and, and so on.

Um, scaling matters a lot, and that, that means that, um, you need compute resources, and you need kind of, um, infrastructure, and you need engineers tab. And then so, uh, um, you know, industry tend to have an advantage when it comes to just, you know, scaling things.

Um, but a lot of the ideas actually came from academia. So, you know, let's take, um, let's take attention, um, which got popular with the transformer in 2017. You know, attention actually was-- has, has been around for a while.

Um, so I think the first mention was in 2014, a paper from Banerdhaw and others and, and Yoshua Bengio, which is coming from, from academia. You know, a lot of ideas d- did come from, from academia. Um, and, uh, scaling things up, of course, has been-- I think OpenAI has been, um, great at, at, um, scaling things up.

Like that's, that, that was the bet that they, they, they made-

Mm-hmm

... you know, after I think GPT-2. So they saw that, oh, scaling Scaling these things up to back then was one point, uh, 1.5 billion parameter. It seemed to give you, uh, amazing capabilities.

Mm-hmm.

So they really committed to that. They really committed to, um, scaling things, and that turned out to be, um, uh-- it's been a pretty successful bet. Um, so, um, I think for, for a-academia, you know, we're still, um, trying to figure out exactly what we're, we're doing, um, in, in this shifting landscape, right?

And, and so, and so lots of folks have been focusing on, for example, evaluation.

Mm-hmm.

Um, so I know, um, the Stanford, um, Center for Foundation Model led by Percy, they have this benchmark called HELM, which is this holistic benchmark. So, um, trying to figure out, okay, characterizing the l- the, the landscape of different kinds of models, what people should evaluate, what people should measure, and things like that.

So evaluation is, is one role. Um, the other one is understanding. So, um, this, uh, this, uh, has, has happened historically where there's been some development in, in the industry, and academia can play a role in explaining, understanding.

Kind of they have the luxury to slow down trying to understand stuff.

Mm-hmm.

Right? So, um, lots of paper on, um, understanding what's really going on, um, probing these models and, and so on, I think. Um, I'm not as familiar with the NLP literature, but my impression is there's, there's a lot of that going on in the, um, kind of NLP, um, conferences, which is understanding what these models are doing-

Mm-hmm

... what capabilities they have and, and so on. Um, and the third one I, I could see is that, um, they-- academia can take, um, more risky bets in the sense that, um, we can work on stuff that, um, quite-- they're quite different from, from industry.

I think industry, my, my impression is, uh, you know, you're, you're trying to-- you have some objective. You're trying to-

Mm-hmm

... say, "Hey, for this quarter, we wanna scale the model in this particular way. Next quarter, we want the model to have these capabilities." Um, and so, um, you're, you're, you're hitting-- you're, you're, you're, you're trying to get objectives that maybe, I don't know, 70% that will, will work out-

Mm-hmm

... and, you know, because it's important for, for the company's direction. Um, I think for academia, um, the mo- uh, the, uh, the, the way things work is like, you know, you have many, many researchers or, or PhD students and, and they're kind of pursuing independent directions, and they have a little bit more flexibility on, "Hey, I'm gonna try out this, you know, seemingly crazy idea, um, and see."

Let's say there's a 30% chance of success or something, right? And however you, you define success.

Mm-hmm. Yeah.

Um, um, for, for academia, a lot of the time success just means like, "Hey, we, we found something interesting." Right? And then, um, that could eventually go into industry through collaboration and, and so on. So I, I do see, um, academia and, and industry kind of playing complementary roles.

Um, and as for s- for someone choosing a, a career, I think, uh, just more, more generally, um, industry would be probably better in terms of compensation-

Mm-hmm

... in terms of, uh, probably work-life balance. Um, and, but my biased perspective is that maybe academia gives you a little bit more, um, freedom to, um, think, um-

Mm-hmm. Yeah

... and understand things. Um, so, you know, it's, it's probably comes down to personal choice. Um, I end up choosing to, um, to be a professor next year at, at Princeton. But of course, like I wanna maintain, um, kind of, um, relationship with industry folks.

I think industry folks can provide very valuable feedback, um, to what we're doing in academia so that we, we understand where the field is, is moving because, uh, you know, some of the directions are very much, uh, uh, influenced by what, for example, OpenAI or Google is, is doing, right?

So we wanna understand where the field is moving, what, uh, what are some promising dir- uh, applications, and try to anticipate, okay, if the field is moving like this, if these are-- these applications are gonna be popular, what problems will be important in two, three years?

Mm-hmm.

Right?

Right.

And then we try to, like, start thinking about those problems so that hopefully in two, three years, like we have some of the answers to, um, some of the, some of these, uh, problems in two, three years. You know, sometimes, sometimes it, it works out, sometimes it doesn't, you know?

Yeah.

But as long as we do interesting things in, in academia, you know, that's, that's the goal.

Yeah. And you mentioned the eval side. So we did a Benchmarks 101 episode, and one of the things we were seeing is like sometimes the benchmarks really influence the model development, you know. Because obviously, if you don't score well on the benchmarks, you're not gonna get published-

Mm-hmm. Yeah

... and you're not gonna-

Yeah

... get funded. Um, how do you think about that? Like, how do you think that's gonna change now that a lot of the applications of these models, again, is in more like narrow industry use cases? Like, do you think the goal of like the academia evals is to like to be very broad, and then industry can do their own evals?

Or what's the relationship there?

Yeah. So I think evaluation is important and often, um, a little bit underrated. So it's not like, you know, as flashy as, uh-

Mm-hmm

... oh, we have, we have a new model that, uh, you know, can do such and such. Um, but, um, I think evaluation, you know, what you don't measure, like you can't make progress on essentially. Um, um, so I think industry folks like, uh, of course, they have specific use cases that, that their models need to do well on, and, you know, that's what they care about.

I think, um, for, um, not just academia, but other groups as, as well, um, people, uh, people do understand what are some of the emerging use cases. So for example, um, you know, now one of the most popular use, use case is, is, um, is chatbot, right?

Mm-hmm.

And then I think folks from, um, um, this organization called, um, from Berkeley, uh, some of them are from Berkeley, called MLSys.

Mm-hmm.

Like they, they set up this kind of chatbot arena-

Mm-hmm

... um, to, to, to essentially benchmark different models. So people do understand what are some of the emerging use cases. People do contribute to like, um, evaluation and, and, and measurement. And a- as a whole, I think people- Try to contribute to the field and move the field forward, albeit at, you know, maybe slightly different directions, but we're making progress and, and definitely, um, evaluation and, and measurement is like one of the-- like the way-- one of the ways you, you, you make progress.

Um, so I think going forward, there's g-still gonna be just, um, more models, more evaluation. We'll just have better understanding of what these, these models are doing and, and what capabilities they have.

Yeah. And I like that your work has been focused on not making benchmarks better, but it's like, let's just make everything faster. Like, let's-- So it's very horizontal. Um, so FlashAttention-2.

Yeah.

You just released that on, on Monday. Um, I read in the, in the blog post that, um, a lot of the work was like also related to like some of the NVIDIA library updates.

Mm-hmm.

Um, yeah, maybe run a bit-- uh, run us through some of those changes and some of the innovations there.

Yeah. Yeah, for sure. So FlashAttention-2 is something I've been work- working on for the past, um, couple months. And, and we've had, um... It actually started, uh, so the, uh, the story is, um, the c- the NVIDIA Cutlass team, um, they released, uh, a new version of their, their library, which contains all these primitives to allow you to do like, you know, matrix multiply or memory loading on, on GPU, uh, efficiently.

So it's, it's a, it's a great library, and I, I, I built on that. Um, so they, they released their version, um, three back in January, and I, I got really excited, and I wanted to play with, um, with that, that library.

Mm-hmm.

So as an excuse, I was just like, "Okay, I'm gonna refactor my code and, and use this library." So that was, that was kind of the start of the, of the project. Um, by the end, and like I just ended up working with the code a whole lot more, and I realized that, hey, there are these, uh, inefficiencies-

Mm-hmm

... still in FlashAttention. Um, we could change this way or that way and make it, in the end, twice as fast. But of course, you know, use, uh, building on the library that the NVIDIA folks, um, released. So that was a kind of a really fun, um, exercise, I would say.

Mm-hmm.

You know, it started out as just like, it's just an excuse for myself to play with- ... with a new library. What ended up was like a month, uh, several months, uh, of like, um, improvement, improving, uh, FlashAttention, discovering new ideas.

Um, and in the end, we, we managed to make it, um, two x faster, and now it's pretty close to, um, probably the efficiency of things like matrix multiply, which probably is the most optimized subroutine-

Mm-hmm

... on the planet. Um, so we're, we're really happy about it. The NVIDIA Cutlass team has been very supportive and, um, yeah, hopefully in the future, we, we're gonna collaborate more.

Yeah.

Yeah.

And since it's a NVIDIA library, can you only run this on like CUDA runtime? So like could you use this and then run it on like a AMD GPU?

Mm-hmm. Yeah. Yeah, so it's an NVIDIA library, so, um, so right now kind of the code we release, um, runs on NVIDIA GPUs, which is, uh-

Mm-hmm

... which is, uh, what most people are using-

Yep

... to train models. Of course, there are emerging, um, uh, other hardware as well. So the AMD folks did implement a version of FlashAttention, um, I think last year as well, and that's, that's also available. Um, the, um, I think there's some implementation on CPU as well.

For example, there's this library, GGML-

Mm-hmm

... where they implemented the same, kind of the same idea running on Mac and, um, and CPU. So I think that kind of broadly the idea would, would apply. Um, the current implementation ended up using NVIDIA's, um, library or, or, or, or primitives.

Um, but I, I expect n-the, the idea to be broadly-- uh, these ideas, uh, to be broadly, uh, applicable to different hardware. As long as-- I think the main idea is you have like asymmetry in, in memory hierarchy, which tend to be everywhere-

Mm-hmm

... you know, in, in a lot of, um, a lot of accelerators.

Yeah. Yeah, it kind of re-reminds me of, um, Sarah Hooker's-

Mm-hmm

... uh, post, like the Hardware Lottery.

Yes.

It's, uh, there could be all these things that are much better, like ar-architectures that are better, but they're not better on NVIDIA, so we're never gonna know if they're actually, um, improved. H-How, how does that play into like some of the research-

Yeah

... that you all do too?

Yeah, so absolutely, yeah. I think, uh, Sarah Hooker, um, the-- she, she wrote this, this piece on Hardware Lottery, and that's, uh, I think she, she captured really well of what, um, a lot of people have been thinking about this.

And, and I certainly think about Hardware Lottery, um, quite a bit given that, you know, I do s- I do some of the work that's kind of really low level at the level of, hey, we're optimizing for GPUs or NVIDIA GPUs and, and, and, uh, optimizing for attention itself.

Mm-hmm.

Um, and at the same time, I also work on, um, other, um, algorithms and, and methods and transformer alternatives. Um, and, and we do see this effect, um, in play, not just hardware lottery, but also kind of software framework, um, lottery.

Mm-hmm.

You know, attention has been, uh, popular for six years now, and so many, um, so many kind of, uh, engineer hours has been spent on making it as easy and efficient as possible to run transformer, right? Uh, there's, uh, there's libraries to do, you know, uh, all kind of tensor parallel, PyTorch parallel-

Mm-hmm

... if you use transformer. Let's say someone else develop alternatives, or let's just take, um, recurrent neurons like LSTM-

Mm-hmm

... GRU, um, right? And if you, if we wanna do that, um, and run that efficiently on current hardware with current software, uh, framework, that's, uh, quite a bit harder.

Mm-hmm.

Um, so in some sense, um, there is this feedback loop where somehow the, uh, model architectures that take advantage of hardware become popular, and the hardware will also kind of evolve to optimize a little bit-

Yeah

... for that kind of architecture. And software framework, um, software frameworks, um, will also evolve to like Optimized for that particular architecture. And right now, transformer is the, the dominant, um, architecture. Um, so, uh, yeah, I, I'm not sure if there is a good way, um, out of this.

Of course, there's a lot of development things like com-- uh, I think compilers, um, will, you know, play a role because compilers allow you to maybe still be much more efficient across different kinds of hardware, um, because essentially you, you write the same code and the compiler will be able to, um, to make it run efficiently on different kinds of hardware.

So for example, there's this, um, language Mojo from-

Mm-hmm.

... uh, Modular, um, AI. You know, they-- they're compiler experts, right? And, and their bet is, um, AI models will run-- will be running on different kinds of devices. So let's make sure that, um, we have compilers, really good compilers with a, a good language that, um, that then, um, the compiler can do a good job optimizing for all kinds of, of devices.

So like that's, that's maybe one way that you can get out of this, this, um, cycle. Um, but yeah, I, I, I'm not sure of a, a good way. You know, in my own research, like I have to think about both the-

Yeah

... kind of new algorithm, new model, and how it maps to hardware. Um, so they're, you know, they're crazy ideas that seem really good, but, uh, will be really, really difficult to run efficiently. And so as a result, um, you know, for example, we can't really scale some of the, some of the architectures up simply because they, they're not hardware friendly.

Mm-hmm.

Um, so I, you know, I ha- I have to think about both, both sides when, when, um, you know, I'm working on, on, on, on new models.

Yeah. Have you spent any time looking at some of the new kinda like AI chips companies, so to speak? Like the Cerebras of the world. Like one of their innovations like, you know, co-locating everything on the chip, so you-

Mm-hmm

... kinda remove some of this like, uh, memory bandwidth issue.

Yeah.

Um, yeah. How do you think about that?

Yeah, I think that's, that's an interesting bet. I think, uh, Tesla also has this, uh, Dojo-

Mm-hmm

... supercomputer, um, where they, uh, they try to have as essentially as fast, um, um, on-chip memory as, as possible and re-re-removing some of these, um, um, uh, data transfer back and forth. Um, I think that's a promising direction.

Um, the, um, the issues I could see, you know, I'm definitely not a hardware expert. Um, one issue is, uh, the on-chip memory tend to be really expensive to manufacture, like much more expensive per, um, per gigabytes-

Mm-hmm

... compared to off-chip memory. Um, so, um, I talked to, uh, you know, some of my friends are at, uh, Cerebras and, you know, they, they have their own, own stack and compiler and, and, and so on, and they can, they can make it work.

Um, the, uh, the other, uh, kind of obstacle is again, with, you know, compiler and, and software framework and, and so on. You know, for example, they can, um, you know, if c- if you can run PyTorch on, on this stuff, like, you know, lots of people will be-

Mm-hmm

... will be using it. But, uh, um, you know, supporting all the, all the operations in, in PyTorch will take, uh, will take a long time to implement. Um, of course, you know, people are, are, are, are working on this.

So I think, yeah, we kinda need these different bets on the, the hardware side as well. Um, hardware has, my understanding is, has a kind of a longer, um, timescale. So you need to design hardware, you need to manufacture it, you know, maybe on the order of three to five years or something like that.

So, um, you know, people are taking different bets, um, but the kind of the, the AI landscape is changing so fast that it's hard to predict, okay, what kind of models will be dominant in, let's say, three or five years.

Um, or thinking back, you know, five years ago, would, would we have known-

Mm-hmm

... that, um, Transformer would've been the, the dominant architecture? Maybe, maybe not.

No.

Right. And so p- different people will make different bets on, on the hardware side.

Yeah. Does the pace of the industry and the research also influence the PhD research itself? So like, for example, in, in your case, you know, you're working on improving attention. It probably took you quite a while to like read the paper and everything.

Yeah.

But like in the meantime, you could have had a new model architecture come out-

Right, right

... and then it's like nobody cares about attention anymore.

Yeah.

Um, how do people balance that?

Yeah. So I think it's, it's tough. It's definitely tough, um, for PhD students, for researchers, um, given that the field is moving, um, really, really fast. Um, I think it, it, you know, comes down to understanding fundamentals. Um, um, because that's, that's essentially, for example, what the PhD allows you to do.

It's been a couple years, um, understanding the, the fundamentals. So for example, when I started my, my PhD, I was working on, um, understanding matrix, uh, vector multiply, which is, you know, uh, is a very-- It's been a concept that's been around for hundreds-

Mm-hmm

... hundreds of years. We were trying to characteri- characterize what kind of matrices would have theoretically fast, uh, multiplication algorithm. Um, that seems to have nothing to do with, you know, AI or anything. Um, but that was, that was a, um-- I think that was a time when kinda I d- I developed kind of, um, mathematical maturity and, and, and research taste and research skill.

Um, and, you know, it doesn't, doesn't-- The, the, the research topic at that point didn't have to be like super trendy or anything as long as I'm developing skills-

Mm-hmm

... as a, as a researcher, I'm, I'm making progress. And eventually, um, you know, uh, I, I've gotten, uh, you know, quite a bit better in terms of like research skills, right? Um, and that allows, um, for example, PhD students later in their, their career to, um, kind of, uh, quickly develop, um, solutions-

Mm-hmm

... to whatever, you know, problems they're, they're facing. Um, so I think that's just the na- natural arc of like, uh, of like how you're being trained as a, as a researcher. Um, for, uh, for a lot of PhD students, I, I think given the, the, the pace is, is so fast maybe It's harder to justify spending a lot of time on the fundamental.

And, and, you know, it's, it's tough. Like, what is-- it's kind of explore, exploit kind of-

Mm-hmm

... a kind of a dilemma, and I don't think there's a, there's a, there's a universal answer. Um, so I personally spend some time doing this kind of exploration, you know, reading random, uh, textbook or, or lecture notes, and I spend some time just keeping up with the latest-

Right

... uh, architecture or methods and, and, and so on. I don't know if there's a right balance. It depends on, um, on, uh, from-- it varies from person to, to person. But if you only spend 100% on one, um, you know, either you only do exploration or only do exploitation, I think it probably won't work in the long term.

It's probably gonna have to be a, a mix, and you have to just experiment and kind of, uh, be introspective and say, "Hey, I, I tried this kind of mixture of, I don't know, one explorer, exploration paper-"

Mm-hmm

"... and one exploitation paper. Like, how did that work out for me? Should I..." You know, having conversation with, with, for example, my advisor about like, "Hey, did that work out? You know, should I shift, I focus more on, on one or the other?"

Like, I think quickly adjusting, that's-

Mm-hmm

... um, and focusing on, on the process, I think that's probably the right way. I don't have, like, a specific recommendation that, hey, you focus, I don't know, 60% on, on lecture notes and 40% on archive papers, anything like that.

Yeah. Um, let's talk about some transformer alternatives.

Mm-hmm.

You know, say Jonathan Frankle loses his bet and-

Yeah

... uh, transformer is not the state-of-the-art architecture.

Mm-hmm.

What are some of the candidates to, to take over?

Yeah. So, so this is a, uh, a, this, uh, this bet is, is quite fun. So this-- I, uh, my understanding is this is a bet between, um, Jonathan Frankle and, uh, Sasha Rush, right?

Mm-hmm.

Um, and, you know, uh, I've, I'd, I've talked to, through Sasha a bunch, and, um, I think he recently gave an excellent tutorial on, um, kind of transformer alternatives as well, so I would, I would, uh, recommend that.

Um, so just to quickly re- uh, you know, quickly, um, recap, I think there's been a l- uh, quite a bit of, uh, of, uh, development more recently about, um, transformer alternatives. So architectures that are not transformer, right?

Um, and, um, the question is, can they do well on, for example, language modeling, which is-

Mm-hmm

... kind of the, uh, the application that a lot of people care about, um, these days. Um, so, um, there are, uh, there are methods based on, um, kind of state space methods like, um, that came out in 2021 from, from Albert Gu and Curran and, and, and, and Chris Rae, um, that are su- you know, presumably could do much better in terms of capturing long-range information, um, while not scaling quadratically.

They, they scale, you know, sub-quadratically in terms of sequence length. So potentially, you could have, uh, a much more efficient architecture when sequence length gets, uh, really long.

Mm-hmm.

Um, the other ones has been focusing more on, um, recurrent neural nets, um, which is, you know, again, an old idea, but, you know, adapting-

Right

... to the, the kind of the new, um, landscape, so things like RWKV. We've also per-- uh, yeah, I've also personally worked on, on, on this, in this space as, as well. Um, so there's been some promising results.

So there's been some results here and there that show that, hey, we, um, these alternatives, either RNN or state space methods, can match the tr- um, you know, the performance of, of transformer on language modeling, so that's really exciting.

And we're starting to understand, um, on the, on the kind of academic research side, we, we understand, like, do we really need attention?

Mm-hmm.

Right? Uh, that's-- I think that's a valuable kind of intellectual, uh, thing to understand. Um, and, um, maybe we do, maybe we, we, we don't. Um, but if we wanna know, we need to spend se-serious effort on, um, on trying the alternatives.

Mm-hmm.

And there's been, uh, folks pushing on this direction. I think RWKV scale up to-- they have a model at fourteen billion that seems pretty competitive with, with transformers, so that's really exciting. Um, um, and, uh, so, you know, that's kind of an intellectual, uh, you know, th-thing.

We wanna figure out if attention is necessary.

Yeah.

Um, so that's one motivation. Um, the other motivation is, um, I think alternatives, uh, uh, transformer alternative could have an advantage, um, in practice, um, in some of the use cases. So, um, one use case is really long sequences.

Mm-hmm.

The other is, um, really high throughput of generation. So for really long sequences, um, when you train with transformer, you know, with FlashAttention and so on, it's still, you know, the computation is still quadratic in the sequence length.

Um, so if your sequence length is on the order of, I don't know, sixteen K, thirty-two K, a hundred K or something, which, you know, some of these models have-

Mm-hmm

... sequence length-

Yep

... a hundred K, um, then you do get significantly, uh, slower in terms of training, uh, also in terms of inference. So maybe these alternative architectures could scale better, um, in terms of sequence length. Um, I haven't seen, uh, you know, actual validation on this as in like, uh, let's say an, an RNN model release with context length, I don't know-

Mm-hmm

... a hundred K or something. I haven't really seen that. But, um, the promise or, or the hope could be that as we scale to long sequences, um, these alternative architecture could be more well-suited.

Mm-hmm.

Um, not just text, but things like high-resolution images, audio, video, and, and so on.

Mm-hmm.

Um, which are, you know, emerging applications. So that's one, uh, long sequences. Number two is a high throughput generation, where I can imagine scenarios where, um, y- the application isn't like a, an interactive chatbot, but let's say a company wants, um, to batch as many requests as, as, as possible on their, on their server, or like they're doing offline processing.

They're-

Mm-hmm

... you know, generating stuff based on their internal documents that you need to p-process in, in batch, right? And, um, the issue with transformer is that, uh- To, during generation, it essentially needs to keep around all the previous-

Mm-hmm. Yep

... uh, history. It's the, called the KV cache. Um, and that could take a significant amount of memory, so you can't really batch, um, too much because you, you run out of memory. Um, for other-- I, I am personally bullish on, on, uh, on RNNs.

I think RNNs, um, they don't-- they essentially summarize the past into a, a, a, uh, a state vector. They have fixed size, so the size doesn't grow with the history. Um, so that means that you don't need as much memory to, uh, keep around all the previous, um, tokens.

Um, and as a result, I think you can scale to much higher batch sizes, and as a result, you can much, make much, um, more efficient use of the, of the GPUs or the, the accelerator, and you could have much higher, um, generation throughput.

Mm-hmm.

Now, this has, I don't think has been validated, uh-

Right

... at scale. So as a, as a researcher, you know, I'm bullish on, on this stuff because I think in the next couple years, like these, these are use cases where these alternatives could have an advantage. Um, we'll just kinda, uh, have to, to, to wait and see, to, to see if these, you know, these, uh, these things, um, will, will happen.

I am personally bullish on, on, on this stuff. At the same time, I also like spend a bunch of time with making attention-

Mm-hmm

... as, as fast as possible. So, uh, I kinda play, I, I, I kinda am ha- maybe hedging. I'm, I'm, I'm playing both sides. Yeah. Um, ultimately, we wanna understand, um, as, as researcher, we wanna understand what works, why do the model, why do the models, um, have these capabilities.

And, and one way is let's push attention, um, to, uh, to be as efficient as possible. On the other hand, let's push other alternatives to be as efficient as, as, uh, we can scale as big as possible, and, and s- and so that we can kind of compare them-

Mm-hmm

... and understand.

Yeah. Awesome. And I think as long as all of this work happens in the open, it's, you know, a net positive for everybody to explore all the-

Right

... all the paths. Uh, yeah, let's talk about open source AI.

Yeah.

Obviously, Together, you know, when Red Pajama came out, which was a, you know, an open clone of like the Llama 1-

Mm-hmm. Yeah

... um, pre-training dataset, it was a, a big thing-

Right

... in the industry. Llama 2 came out on Tuesday? I forget.

Yeah. Yeah.

Man, this week-

Yeah

... there's been a lot of things going on.

Yeah, yeah, yeah.

Um, which, you know, they call open source, but it's not really open source.

Mm-hmm.

Um, actually, we wrote a post about it that was on the front page of Hacker News before this podcast.

Okay.

I was frantically responding.

Mm-hmm.

How do you think about what open source AI really is? You know, like in my mind, there's, in open source software, we have different levels of open.

Mm-hmm.

So there's like free software that's like the GPL license.

Mm-hmm.

There's open source, which is, uh, Apache, MIT.

Yeah.

And then there's kinda like restricted open source-

Mm-hmm

... which is the SSPL and some of these other licenses.

Yeah.

In AI, you have the open models. So Red Pajama is an open model because you have the pre-training dataset, you have the training runs-

Right

... and everything.

Right.

Uh, and then there's obviously Random Less that doesn't make it one-to-one if you retrain it.

Mm-hmm.

Then you have the open weights model-

Yeah

... that's kinda like, um, Stable LM-

Mm-hmm

... where the weights are open, but the dataset is not open.

Right.

And then you have, um, Llama 2, which is the dataset is not open. The weights are restricted. It's kinda like not really open source-

Mm-hmm

... you know, but, uh, open enough.

Yeah.

I think it's net positive because it's like $3 million of flops-

Mm-hmm

... like donated to the public.

Yeah.

You know?

Yeah, yeah.

Um, how do you think about that? And also like, uh, as you work with Together, you know-

Mm-hmm

... what is, uh, your philosophy with open source AI?

Right. Right. Yeah. That, I think that that's, that's a g-r-great question, and, um, I, I, I, I, I think about it on maybe more practical, um, terms. So, um, you know, uh, of course, like, uh, Meta has done an amazing job training Llama 1, Llama 2, and for Llama 2, they, um, kind of make it much less restrictive compa- compared to Llama 1's, where, um, now, um, you can use it for businesses unless you are, however, you know-

700 million, uh, monthly active users

... monthly active user or something like that. Um, I, I, I, I think just this change will have a very significant, um, impact in the kind of landscape of, of open source, um, AI, where now lots of businesses, lots of companies will be using, I, I, I, I expect will be using things like Llama 2.

They will fine-tune on their own dataset. They will be serving, um, you know, variants or derivatives of, of Llama 2. Um, whereas before, you know, with Llama 1, it was a als-

Mm-hmm

... so really good model, um, but your business, uh, companies weren't allowed to do that. Um, so I think on more practical term, it's kind of shifting the balance between kinda closed source model like OpenAI and, and, and Anthropic and Google, where you're making API calls, right?

And you, maybe you don't understand as much of like what the h- what the model is doing, um, how the model is changing and, and, and so on. Versus now we have a, a, you know, an, uh, a, a model with open weight-

Mm-hmm

... that is, uh, you know, pretty competitive from, uh, from what I've seen in terms of benchmarks, pretty competitive with GPT 3.5, right? And if you fine-tune it on your own data, maybe it's more-

Mm-hmm

... more well-suited for your own data. Um, and I, I, I do see that's gonna shift the balance of like more and more folks are gonna be using, um, let's say derivatives of Llama 2. More and more folks are gonna fine-tune and serve their own model instead of calling an API.

Um, so I think that, that shifting of balance is important because we don't, and, you know, in one way, we don't want a, just a concentration of decision-making power, um, in the hands of, you know-

Mm-hmm

... a few, a few companies. Um, so I think that's a really positive, um, development from, from Meta. Of course, you know, training the model takes a couple millions of dollars, but like, you know, engineers have, and they've, they've-

Mm-hmm. Right

... I'm sure they, you know, they spend tons of time trying to, trying many, many different things. So the, the, the actual cost is probably way more than that, and, and they're releasing it in the, uh, you know, they make the weights available and, and, um, they allow And probably a lot of companies are gonna be using this.

Mm-hmm.

So I think that's a really positive development. Um, and we've also seen amazing progress on the, you know, open source community where they would take these models and they either fine-tune like, um, on, on, on different kinds of datasets or even make changes to the, to the model.

So as an example, um, I think for LLAMA1, the context length was limited to 2K.

Mm-hmm.

But like a bunch of folks figured out some really simple methods to scale up to like 8K.

Yeah, like the rope-

Yes

... uh, thing. Yeah.

Yeah. Yeah. So, um, I think the open source community is like very creative, right? And, you know, lots of people. Um, so, uh, L- LLAMA2 will, uh, again kind of accelerate this, where more people will try it out, more people will make tweaks to it and make a contribution, and then so on.

So overall, I think I, I see that as still a, a very positive development for, for the field. Um, and, uh, there's been, um, lots of libraries now that-- or, or, um, libraries that will allow you to host or fine-tune these, these models, like even with quantization-

Mm-hmm

... and, and so on. Yeah, just a couple hours, uh, after LLAMA2 was released, like tons of, um, companies annou- uh, uh, announcing that, "Hey, it's on our API or hosting," and then so on, and, and Together did, did the same.

Um, so it's a very fast-paced development, and just having an, um, just kind of a, a model with the, with available weights that business are allowed to use, I think that alone is already very, um, positive development. But at the same time, yeah, we, we can do, you know, much better, um, in terms of releasing dataset.

I think dataset tend to be, um, somehow people are not incentivized to release dataset.

Right. Yeah, yeah, yeah.

Um, so, you know, philosophically, yeah, you wanna be as open as possible, but on practical term, I think it's a little bit harder for companies to release dataset.

Mm-hmm.

Um, you know, legal issues. Um, uh, the dataset released tend to be not as, uh, kind of, uh, eye-catchy-

Mm-hmm. Yep

... as the model release, so maybe people are less incentivized to do that. Um, we've seen some, you know, uh, quite a few companies releasing dataset, you know, Together, um, released, uh, RedPajama, uh, dataset. I think Cerebras then worked on that and, you know, deduplicate and cleaned it up and released Slim Pajama and, and so on.

Mm-hmm.

So we're also seeing positive development on that front, kind of on the pre-training dataset. So I, I do expect that to, um, to, to continue. And then on the fine-tuning dataset or instruction fun- tuning dataset, I think we now have quite a few open datasets on, on instruction tuning and, and fine-tuning.

Um, but yeah, these companies still, they, they do pay, uh, for human labelers-

Mm-hmm

... right, to, to annotate these, uh, instruction tuning dataset, and that is expensive.

Right.

Um, and maybe, you know, they will see that as their competitive advantage. Um, and so it's harder to incentivize these companies to, to release, um, these dataset. So I think on, on practical term, we're still gonna make pro-- a lot of progress in, on open source AI, on both the model development, on both model hosting, on both, uh, on, on, on, um, pre-training dataset and, and fine-tuning dataset.

Um, you know, right now maybe we don't have kind of the perfect, uh, like open source model-

Mm-hmm

... since the, oh, the weights are open, uh, weights are available, uh, all the datasets are available. Maybe we don't have such a thing yet, but we've seen, um, very fast development on the, on the open source side right now.

I think just maybe this time last year, there weren't as many models-

Right. Yeah, yeah

... as are competitive with, let's say, you know, ChatGPT.

Mm-hmm. Yeah. Yeah, I think the open datasets, they have so much more impact, you know, than open models. If you think about Elluther-

Mm-hmm. Yeah

... and like the work that they've done, GPT-J was like great, and like the Pythia models are great, but like the Pile and like the Stack are like-

Absolutely. Yes

... you know, everybody uses them.

Mm-hmm.

You know? So hopefully we get more, more people to contribute time to work on datasets, you know, instead of doing the 100th, uh, open model that like performs worse than all the other one, but they wanna say they released the model.

Yeah. Yeah.

Um-

I think, you know, we-- maybe like the, the, the question is how do we figure out a kind of incentive structure so that companies are willing to, uh, release, uh, open, uh, datasets and, um, and, um, so, you know, for example, it could be like I think some of the, uh, some of the organizations are now doing this where they are, um, kind of asking volunteers to like, you know, annotate and, and, and so on.

Mm-hmm.

And then kind of maybe the Wikipedia model-

Right

... of, of like dataset or especially for sh- instruction tuning could be, could be interesting, where people actually volunteer their time and, and instead of editing Wikipedia-

Mm-hmm

... like, you know, add annotation. And somehow they're, they acknowledge and, and, and feel incentivized to do so. Um, hopefully we get to that, that kind of level of in terms of data would be like kind of like, uh, Wikipedia.

Mm-hmm.

And in terms of model development, it's kinda like Linux where people are contributing patches and, and improving the model in some way. Um, I don't know exactly how that's gonna happen, but, um, uh, based on history, I think there is a way to get there.

Yeah. I think the DALL-E 15K dataset is like a good example-

Yes. Yeah

... of a company saying, "Hey, let's do this smaller thing."

Mm-hmm.

"Just make sure we make it open."

Mm-hmm.

Um, yeah. And it came out very-- We had Mike Conover from Anthropic-

Oh, nice

... on the podcast, and he was like, "People just bought into it, and like leadership was bought into it."

Mm-hmm.

You know, you have companies out there with like, you know, 200,000, 300,000 employees.

Yeah.

It's like, just put some of them to label some data, you know? Like it's gonna be helpful. Um, so I'm curious to see how that evolves. What made you decide to join Together?

Yeah. So, um- Uh, for Together, yeah, the, the focus has been focusing a lot on open source model, and, and I think that aligns quite well with, with, um, what I, I care about. Of course, um, I also know a bunch of people there that, and I, that I, I know and trust, and I'm excited to, to, to work with them.

Um, philosophically, uh, I think the, the, the way they've been really open with like dataset and, and model release, uh, I, I, I like that a lot. Um, personally, I've, you know, uh, for, for the stuff, for example, the research that I've developed, like we try-- we also try to make, you know, code available free to use and modify and, and-

Mm-hmm

... and so on, con-contributing to the, to the community. And, you know, that has allow-- has given us really valuable feedback from the community in improving our, our, our work. So, um, philosophically, I like the way, uh, you-- T-Together has been focusing on open source, um, open source model.

Um, and the nice thing is, um, we're also gonna be kinda at the forefront of, of research and, uh, and the, uh, the kind of research areas that I'm really excited about, things like efficient training and-

Mm-hmm

... inference, um, aligns quite well with, uh, what the company is, is doing. We'll, we'll try our best to make things open and available to, to everyone. Um, yeah, but it's, it's gonna be fun being at, at the, the company doing, you know, leading a team doing research on the topic that I, I, I really care about.

And, and hopefully, we'll, we'll make things open to, to, um, to benefit kind of community. Yeah.

Awesome. Um, let's jump into the lightning round.

Okay.

We usually have three questions. So one is on acceleration, one on exploration-

Okay

... and then a takeaway. So the first one is, what's something that already happened in AI machine learning that you thought would take much longer than it has?

I think, um, understanding jokes.

Mm-hmm.

I, I didn't expect, uh, that to happen. Uh, but, uh, you know, it turns out scaling model up and training lots of data, um, the model can now understand jokes. I, uh, maybe it's a, it's a small, small thing, but like that was amazing to me.

What about, um, the exploration side? What are some of the most interesting unsolved questions in the space?

Um, I would say reasoning. Um, you know, in, in a broad term. Um, we don't really know how these models are do-- Essentially, you know, they do something that looks like reasoning. We don't know how they're doing it.

Um, we have some ideas, and in the future, I think we can-- we, we will need to design architecture that kind of explicitly have some kind of reasoning module in it, um, if we wanna, uh, have much more capable models.

Mm-hmm. Um, what's one message you want everyone to remember today?

Um, I would say try and let's-- try to understand both the algorithm and the systems that these algorithms, um, run on. I, I think the-- at the intersection of machine learning system has been, um, really exciting, and there's been, um, a lot of amazing results at this intersection.

And then when you, you know, scale models to, to large scale, um, both the machine learning side and the system side really matter.

Awesome. Well, thank you so much for coming on Tri.

Yeah.

This was great.

Yeah, this has been really fun.

Awesome.