#324 Sharon Zhou: Inside AMD's Plan to Build Self-Improving AI

0:00Catastrophic forgetting is definitely a problem, especially in post-training and when you don't have access to the original pre-training data.

0:08Sharon Zhou:How much of that is focused on improving the design of AMD hardware and how much of it is focused on general model development, which AMD is not doing, right? AMD does not develop its own models. I think people want infinite chips, Craig, so no, it doesn't relieve the pressure. I am Sharon. I'm the VP of AI at AMD. And I think about self-improving AI, self-improving LLMs, which we'll get into later. But my background comes from AI research. So I used to be an AI researcher at Stanford, where I did my PhD with Andrew Ng. I taught there as adjunct in generative AI back before all this chat GPT stuff.

0:52And after Stanford, I started a startup, an AI infrastructure for startup doing post training of language models actually on AMD GPUs. This was started a couple of months before chat GPT launched. And most recently over the last several months, we have transitioned now to AMD. So my team and I are there now and very excited to enable more people to use compute and to get access to compute because that really is the limiting factor. one of the big limiting factors for developing AI and being able to enable more people to steer these models. So that's what I'm really excited about. And yeah, that's why I'm here.

1:36Sharon Zhou:Yeah. And I do want to talk about self-improving AI. Can you start by defining what we're talking about when we talk about self-improvement? Are we talking about models that rewrite their own code or something like refining their own training data? Yeah, I think that's exactly it. It is a broad category, but essentially it's the idea of these models being able to edit any part of themselves to improve themselves. So whether that be the data, whether that be the actual model architecture, whether that be how they evaluate themselves. And actually the part that I'm working on is below all of that, and it's actually how fast they actually run on the GPUs themselves.

2:23So they are writing the kernel code that underlies these models to run faster on these GPUs and to run effectively on them. Yeah. And on new hardware, too. So that's been really exciting to see. Yeah.

2:38Sharon Zhou:I just read about kernel evolve. Yes. Who is doing kernel evolve? I've forgotten. Is that Google? Meta. Is that your guy's work? So there's a lot of different pieces of work around kernel generation and being able to use LLMs to generate these kernels. So we're doing some of that, but I think that work might have been from Meta, but there's some work across. Yeah, I think there's work across the board, across the industry that is very important, I think, towards this end because it enables more people to get on different types of compute. What we did most recently was in collaboration with actually a bunch of different institutions like Meta, Google DeepMind, ML Common, Stanford, NVIDIA, et cetera, was a NeurIPS tutorial on generating kernels using AI.

3:34And so we presented that and basically goes through how we're using AI agents to generate these kernels and how we're thinking about post-training these models to generate kernels more effectively. because what's really exciting about kernel generation and kernel development is actually we have the profiler. So we have the ability to actually see how fast the generated kernels are on the chip themselves. And so that's really exciting. My team, we're also working on a kind of more robust production level benchmark to share with the community as well as different techniques to modify the models to do better on this task.

4:16Sharon Zhou:Okay. Yeah. And for, again, I'm trying to make my podcast a little more accessible. Oh, yeah, of course. So kernel in this sense, because kernel is used in... Everywhere. Too many places. But kernel in this sense, you're talking about a small piece of software that lives on the processing unit that performs a specific task. Is that right? And what are the kernels that you're generating from AI? So there are many layers of the stack for AI. there is the model layer where we kind of build out the model architecture. This is where people talk about transformers and attention and kind of write that out in PyTorch or Jax or TensorFlow.

5:13And so people are like building out models there and then they're using maybe hugging face on top or they're using different tools to leverage those models. Now underneath all of that are ways where the model is running on the GPU, and there are many, many layers, but one of the layers is making those models run really fast on the GPU. And sometimes you can break that down into small different pieces because it's a lot of matrix multiplications, for example. It's a lot of different operations happening on the GPU. And so for each of those operations, you can have a way of basically optimizing the speed of that on the GPU itself.

5:56And so for a given piece of hardware, it may have been designed to be of a certain speed, but you do have to write the software to connect the models of today and of tomorrow, ideally, to that hardware. And that's kind of what that layer is doing. It's improving the efficiency. So it's utilizing the GPU effectively, both the memory. So like what the GPU has as memory and storage, as well as the actual raw compute power. So being able to schedule all of that and use that effectively and maximally is what kernels do. And this is really important because it's expensive, right? Like GPUs over a long time and over just a large capacity when you're parallelizing, it becomes expensive.

6:44And so you want to be able to eke out as much as you can on each GPU. Yeah.

6:51Sharon Zhou:And I understand how. And you're using evolutionary strategies to develop these kernels? Yes. I would say it's a combination of evolutionary strategies, agentic strategies, as well as different types of post-training strategies to be able to get these models to actually be able to write that code effectively, or at least assist our internal kernel engineers to do so. And I think one of the most important things is to do so on useful kernels as well, the most important kernels for language models. Yeah. Well, again, for listeners, what are the most useful? What's what's an example? Well, I would say matrix multiplication.

7:36So when you're doing, when you're learning about the math of a neural network, you'll see that there's a lot of matrix multiplications and a lot of those need to be optimized. And, you know, at first you might think, oh, isn't there just one, you know, two matrices multiplying, but actually no, you can optimize different size matrices that multiply together and you can get like even better performance for something of a different size. Like you can optimize that. So there are thousands, possibly hundreds of thousands of, of those kernels that you could optimize, um, for, for language models. And of course, for a particular, um, customer, particular, like foundation model company, um, they may have a certain set that they particularly care about, um, for their architecture.

8:20And, um, those are ones that are, are high priority. Yeah.

8:25Sharon Zhou:And these kernels in the past have been written manually. Is that right? So having there's a tremendous productivity gain and having AI either do it autonomously or do it as an assistant. What sort of impact does that have on AMD? I think it's having larger and larger impact. So I would say that, yeah, so historically, they've been written manually. And just to give you a sense of how much knowledge this kernel engineer needs to have in their head, they need to know about the GPU architecture. So not just general GPU architecture, not just your average architecture class, but like that generation of the GPU and what changes have been made, you know, what is actually, you know, how to actually write the code to run the software on it and all the different, you know, possible new things that have been invented at the hardware layer.

9:25They have to understand all those things. And then they also need to understand at least enough about what this matrix multiplication is doing, right? And to actually then write that optimally in the code. And I would say that having those two areas of expertise in one person's head is quite rare. And as a result, it's very, very valuable, but also it is bottlenecked at a lot of hardware companies or at actually a lot of different companies, including Frontier Labs. They also have people writing kernels to speed up their models as well, since they have more visibility on their own models that may not be fully visible to everyone else.

10:05And so I think it's a rare skill. And if you're listening, I encourage you to go learn about it. But of course, we're also teaching the models about this. Right. And so what's exciting is for AMD, we have the Rockham stack, which is open source. That's the equivalent of NVIDIA's CUDA stack, which is not as open. And so that being open actually helps us from a language model perspective because the models can kind of read that data, right? And train from that data and use that data to then learn about the Rockham stack and then write those kernels for us. And so that's what we're doing here. And that's the impact is massive from both a, I would say, like direct customer lens of what it could do to improve that.

10:50But also, I would say, and the long tail. So there are researchers working on various different models that are not just Chachi PT or, you know, just Grok or just, you know, just a few of these frontier models. But they're working on a bunch of different other models and they're creating, they're inventing the next model. Right. And those also need to be optimized. And I think having something that is almost just in time, ideally just in time, but, you know, like almost automatically give you a optimized kernel for is very essential to enabling just the entire ecosystem to be able to move over.

11:31Yeah.

11:31Sharon Zhou:And when you say almost just in time, is this is being done on the fly or? That is my goal, Craig. That is my goal. Yeah, because you talk about self-improving AI, that suggests that there are models that are evolving over time or improving over time. And that leads to the question of continual learning, which is kind of the holy grail. but there's catastrophic forgetting and expandability of neural nets and all those problems. Can you talk about where the research stands on continual learning? Yeah. So catastrophic forgetting is definitely a problem, especially in post-training and when you don't have access to the original pre-training data.

12:32Because I think we found in the literature at least that even if you include, originally it was like if you include 20%, but actually if you include only 1 % of the pre-training data back in during post-training, you can actually prevent catastrophic forgetting, basically enabling the model to actually like connect back to its representations back in pre-training, or at least help significantly for that. And I think it, again, depends on data access in terms of who is doing some of that post-training. But I think that's what the literature points to. So it's the extent to which you can access some of that data or use some of those online data sets to just bring back some of those data points from retraining that could actually help your model significantly in preventing catastrophic something like catastrophic forgetting.

13:24But of course, I think this is something that needs to be continually monitored. So I think the frontier labs maybe don't have this problem as much, But as an individual who is probably doing some type of post training, that is something that you do have to consider, especially as you get your especially as your workload gets heavier and heavier. Small little bits of fine tuning that doesn't really change as much.

13:46Sharon Zhou:But so self-improving, are you? Yeah, define self-improving for me. I mean, what are you improving in or attempting to improve in a model? So the part of self-improving that we've been working on is using these language models to write low-level kernel code. So this is code that makes the language model itself run really fast on AMD GPUs. So get the model to actually write its own code to run even faster on the GPU so that it can learn even faster. And so this is the self-improving loop that we're enabling the models to do here. That's really exciting because then that'll enable more people to access this compute and it'll enable more powerful models within a shorter amount of time.

14:33Sharon Zhou:And is that something then that you deploy with AMD hardware that people can use with their models? Or is it a separate model that you make available to users of the hardware? Yeah, this is such a good question. Right now, it's being used internally, and we're releasing different data sets and evaluation benchmarks and methods out into the world. So we did a Nurex tutorial that presented a lot of this work, and we're preparing for ways to actually package this and make it possible for our customers to use it as well. Yeah. And then the idea of editing a model, editing its own training data, or why would you want to do that?

15:27Sharon Zhou:And can you give me an example? And how is that being done? Yeah, so we're editing the kernel code for the model. In terms of editing training data, what that looks like is usually this is like a synthetic data setup. So the model is generating data for the next round of training. So yeah, self-improving AI is more general than just having the models generate kernels and write these low-level kernels to make the models really fast on GPUs. It could also touch on generating data for itself, as well as generating evaluation for itself, and also generating different architectures for itself as well.

16:12So basically using the models to improve itself at any stage of the pipeline. I think data generation is really interesting, just synthetic data generation overall. And there's been a lot of work and research and speculation on it. And I think, you know, we have found that using synthetic data to train these models is very helpful, especially if it's actually not a frontier model and you're distilling from a frontier model. It's not quite distillation, sorry, but it is using synthetic samples from a frontier model to then supervise a smaller model that you might be training. That's very effective.

16:48And I think like there's also debate around synthetic data generation and how much that may or may not lead to the collapse as the model trains since it's looking at its own kind of data. But I do think it is an important space to continue monitoring and to continue understanding. Yeah, so I think what's really interesting is this holy grail that everyone's looking towards of having the model fully improve itself and create its own next generation. I think we're both further than we think and closer than we think. And so I think we're further than we think in the sense that there is still a lot of human expertise that goes into how these things are laid out and what we should try next, for example, of what we should add into the model, whether it be new data source or not.

17:43But I do think there's more than we think as well in terms of the use of AI to help us code in general. So the amount of code that AI is writing is pretty enormous, and it is supporting a lot of engineers that are putting these together and creating the next generation of that AI. So I think it's both we're further and closer than we think. Basically, we're further than we think from the like ultimate like okay there's only one person works at open ai now it's only sam there and it's just the end with chat gbt versus um uh we're we're closer than we think because actually every everyone is actually using um uh chat gbt quite significantly um to write that code so i i think uh it's it's in in in that state right now um but it's it's not exactly like autonomously improving itself fully.

18:40There are pieces that you can launch that feel more autonomous, where you could have it like reflect on its own output and then improve that. It can look at the profiler information so it can look at how fast a kernel ran on the GPU and improve that. But it's not quite like, I'm going to let this go and come back in a month and it will be fully solved.

19:04Sharon Zhou:yeah uh yeah well let's can we back up a little bit uh as i said at the beginning amd's a hardware company you're uh in charge of ai at uh at amd part of it yeah part of it yeah is is that uh uh is is that a research function um because uh and and how much of that is focused on improving the the design of of amd hardware and how much of it is focused on uh on general model development uh which which amd is not doing right amd does not develop its own models oh we have actually uh pre-trained some models and um released them but they are small in the um single digit billion parameter range um and yeah we're probably like not developing models in the sense of like not chat gpt um but uh but yeah we we do have some um that are uh open source models um yeah so i focus on i almost call it like product research i think there is a research element to it that we don't know how good um the models are at solving this full task and end um but there is a product component in that this is part of the loop with what becomes customer facing right so then this becomes part of the loop of what we um want to make available to our customers and so uh it's it's shippable and there is a production component to it.

20:43And so I would say it kind of sits at that intersection, which I know a lot of things today sit at. So, yeah, I think that's one of my teams works on that. And the other team works on, I would say like research and education. So we're thinking about how to stay at the forefront of research and engage with researchers everywhere, but also educate the world and evangelize, you know, educate the world about AI and evangelize AMD GPUs that are able to support all of those new AI workloads. And so, yeah, that's a couple teams that I've been very grateful to be leading. Yeah.

21:27Sharon Zhou:And we tried to connect at NeurIPS, and I'm sorry that we weren't able to. And I missed the talk. can you go through what the presentation that you gave at NeurIPS? Yes. So it was a tutorial, two and a half hours. And my team did a huge part of the presentation. So definitely not just me here. And it was a collaboration across a lot of different institutions. So like, you know, Stanford, Google DeepMind, Meta, NVIDIA, ARM, et cetera. And so it was a really, really great group of people. And what we presented in our tutorial was a few different things. First, because it's a NERVUPS audience, AR researchers, they might not be actually as familiar on GPUs and hardware and how the GPU actually works.

22:19And so we kind of stepped through in hopefully very easy to understand ways, the actual GPU architecture and what pieces that an AI researcher might find interesting, like the memory and like how things are scheduled in. Like something I find interesting is, you know, when you change your batch size, why suddenly like your your job takes forever, right? Your training job takes forever. Oh, it's because of, you know, X, Y, Z happening in in the actual GPU. you. And so that's what we go through first. And then we go through kind of different kernels. So like showing you what a kernel actually is in code in a very simple kernel, and then showing you what it looks like when it's optimized and unoptimized.

23:04So you can understand like, oh, when a kernel runs fast or slow, what it looks like. So if something runs really slowly, that means like ChatGPG could answer you in like minutes as opposed to seconds, right? So that's ultimately what it feels like. But then when you go down into like just a single matrix multiplication, what does that look like? Or a single addition, you know, like X plus Y. What does that look like in terms of time and how do you profile it? And then we go into, okay, what are AI agents and methods that we're using for these self-improving AI essentially? But essentially what are some of these methods that we're using and sharing a lot of that.

23:44So like what kind of things can you do with agents? You can have them, like you can write me a more optimized kernel. Okay, it does it. It maybe is or isn't optimized. You have to check the correctness of it. Maybe it's cheating. And then when you profile it, you can actually get the answer back. And you're like, actually, it was only 1.2x faster. I want it to be even faster. and then you can continually evolve it. That was based on Google's alpha evolve paper, but essentially like continually improve it after multiple calls. The profiling information also provides a really interesting way of collecting or creating an RL environment so that these models can learn from essentially a verifiable reward from the profiler itself.

24:32So it's able to give this number back as, okay, this is how fast your kernel was that you generated and tells the model in a very verifiable way. Kind of like when Chachapiti, they were training on math tasks. It was very verifiable whether the math was correct or not. And this helps the model improve in post-training. And so, sorry, I'm trying to make this accessible, but like, so basically we go through all of those different techniques as well and we share that with the community. Yeah.

25:04Sharon Zhou:And this is done at what point in the process? Someone, a customer is using AMD hardware and they have a model that they want to run on it. And so they're optimizing the kernel before deploying the model. Or is it after deployment and they want to improve it? Yeah. So I think it's a combination of those things. So I think you could use it definitely before you deploy something. So before you deploy something, you're like, okay, this isn't fast enough given my budget with this many GPUs and we need it to hit these markers. Right. And so because that basically translates to TCO translates to the total cost of ownership.

25:50And so let me reduce the speed of all, sorry, yeah, let me make all of these faster inside of the model so that the full model runs really quickly. So that could be before deploying something. It could also happen afterwards, too. A lot of different kernels are written afterwards as well to just speed things up even further. So it could be something you've already deployed, and now you want to speed things up to make use of the hardware you have even more. So I think it could occur in both of those settings.

26:21Sharon Zhou:Yeah. And this is, in my mind, this is more automation than self-improvement, right? You're automating the kernel writing process, and that improves the model. But where does self-improvement come in? Yeah, so the self-improvement is, so I view it as self-improvement, not automation, but I view it as self-improvement because the kernel itself is what the model is running on. So it is the model's own code and it enables the model to run even faster on the GPU. But I guess when broken up, it could be seen as just like, oh, automation. And I do think, I honestly think everything is technically automation.

27:11But if we want to view it through a lens that's like more exciting, where it feels like autonomous, that is just a different perspective on it. Yeah.