llm.c's Origin and the Future of LLM Compilers - Andrej Karpathy at CUDA MODE

Key Takeaways

to introduce to you a legendary person someone who has been hacking and educating at the Forefront of AI for over a decade from neuron networks to computer vision from lalar processing to reinforcement learning he has pushed the boundaries and Inspire Millions all over the world including I think all of us here he is a distinguished machine learning Superstar a founding member of AI the reference human for image net and X Google brain X Deep Mind ex Tesla Mr autopilot he has really seen it over and some month ago on a memorable day this special person joined the kudam mode Discord to start hacking with others on llmc which became one of the greatest and most active Community projects in our server but I guess it's best if he if he tells the story himself so please join me in welcoming The Incredible One and Only Andre kathi wow [Laughter] okay okay yeah I'm very excited to be here this is my favorite kind of event to present that so uh yeah thank you for the invitation and thank you for running Cuda mode and putting this on this is a wonderful event okay so I'll tell you a bit about L.C so what are we doing we're training Transformers uh in C and a pinch of C++ uh how do I do next here okay so I'd like to tell a story a little bit of how this project came about and what this looked like from my perspective so roughly a year ago I was trying to add a video to my YouTube series and I was trying to teach people um LM training gbd training and so on and I was basically hacking on them GB it to work so uh that was me and then um you you've all worked with py of course right so the trickiness comes that okay you have your model which you've written that makes sense but now you have to keep track of a number of abstractions here at the same time so you have to put it to a device you want to compile it you want to wrap it in DDP and uh suddenly things start to be a little bit more complicated because I'm not even sure like in what order do you do these what exactly happens what are these abstractions what do they do to your model so I don't fully understand how any of this works and then what happens is um you want to use your model in different ways so you want to use it in evaluation in training or model inference and so on and what happen to me is that I was able to train the model but for some reason eval and inference was not working and what happened was I was getting some kind of a torch compile eror when I was trying to run my eval and my inference and this is just an illustrative example of a torch compile error it was something else I don't remember I didn't capture it but both of them were given me error inference and eval and a different error and I had no idea what was going on so I did what anyone would do in my position I went to CH discuss I'm looking for PG BLK to solve my issue unfortunately PG BLK did not have any guidance that I could see on that specific error so I was kind of stuck honestly so uh two hours later of fighting with torch compile and trying to figure out what the hell was going on I'm kind of a sad and out I don't know exactly how to solve this and so I felt like I was going to in the beginning I was in denial I was like this can't be happening to me I'm not doing anything crazy I'm just training a little GPT like why is this not working this seems really simple I'm not doing anything crazy and then eventually I entered the stage of anger and I was like okay you know going I write the whole thing like I understand in my mind what I'm trying to do like the computation itself the algorithm itself is totally clear in my mind and for some reason Tor compile doesn't let me like use it run it Etc so I felt a little bit powerless and I was like okay I'm going to take life into my own hands and being control of my destiny I'm going to just write this and see how B could be so uh let's think about like what is p offering you really and there's many things but maybe some of the things that are relevant here I don't know why those bullet points are I don't know what on my slides it's totally fine so I don't know okay but number one we're getting in Array right so a very useful end dimensional array that we can manipulate with operations if we're going to abandon this then we're going to have to do a lot of pter Matic basically making sure that we Ravel and unravel indices correctly uh second we're getting autograph for free so if we don't have autograph we need to do forward and backward passes of all the layers uh we don't have device so we have to worry about memory being on the host or on the device and shoveling memory around your different uh devices between CPU and GPU and so on we don't have simple DP conversions so we have to be very mindful what tensors are stored in what precisions and convert explic between them we don't have t compile so we're going to have to do all the kernel fusions that we want manually and we're going to have to optimize for space Time Performance manually uh and finally we don't have distributed so we're going to have to manually spin up all of our processes make sure that they can find each other communicate with nickel Etc so pyro is really really nice and this is just some of the things that PTO offers so without pyto we're kind of naked in the world right but maybe it's okay so yeah how how bad can it be so step one we have our pyro code which Now isn't the primary thing we're working with it's only a reference that we check correctness with respect to and so we're in P torch land everything is nice and clean we have a little Transformer a few modules and we're just calling them so everything is great and that now becomes our reference in pure um I'd like to just take you through one example of a layer so for example layer Norm here is like a p layer and uh we'd like to basically P this over to see so what kind of process that we go through well we're going to iterate through all the layers number one we need the forward pass and actually I had to write the forward pass of layer Norm um because uh pytorch doesn't just have this kind of implementation in pytorch of layer Norm because it's kind of like a block and eventually PS into some Cuda kernels uh so I had to write a forward passive layer norm and make sure it's equivalent to the layer norm inro and then of course I had to write the backward pass of layer Norm this is where you kind of take out your pen and paper do some backrop uh this is for batch Norm but layer Norm would be similar and yeah we have to write the backward pass and again you're this is all still in pych but it's explicit and you're just making sure the layer Norm of pych P backward matches this represent this uh um basically um manual uh tensor bed um implantation so now we have P code forward backward so next thing we do is we try to Port it to C and this is actually like a lot simp in many cases that you might think so on the left we have the P code and on the right we basically have the equivalent C no forward in C and it's not that crazy right so uh unlike in py we just have a bunch of uh float star arrays uh so we have a float star out float star inputs outputs means standard deviations weights and biases and some high parameters and one thing I really likeed to do in that I just want to keep things simple I don't want to create a tensor abstraction I don't want to create any abstraction really it's just float arrays and operations on float arrays like why should it be a lot more complicated than that so everything is just float arrays everything is fully self-contained there's no underlying representations abstractions to call import Etc this is the lay forward on float arrays and that's it so that's the forward um and then you also do the backward for all the layers once we've done that for all the layers and converted everything to C and make sure that everything matches our reference implementation we have to start to string it together so uh we go into our C code in Main and we have to allocate all of the memory that we're going to be using in L.C all of the allocation happens a single time at the beginning so we pre-plan all the memory that we're going to ever use then it's fixed and from then on it's just dynamics of just feeding data through it and training the model so we have to pre-plan all the tensors their sizes and we have to do that for the parameters and we have the data grad and the m and b for the ADM W buffers and then for the activations as well and we need space for both data and grab and so you just pre-plan all the memory you allocate all of it and then we need to stitch it all up so uh we have all these layers and they have a forward and a backward pass in back ration and so on the forward uh the forward pass just kind of like you allocate all these tensors and you're very careful and you index into them properly and you make sure everything flows correctly through and you just go all the forwards and then all the backwards and then uh you're kind of done and you're left with gradient and can do an update so stringing that together is the second piece of work and then once we sort of like Str it together you get something that you can just compile and run so uh we um on the top left is everything that's required we download a star pack which is really just the gpt2 weights in a single binary file very simple uh and also we need the data set in this case Dy Shakespeare and the tokenizer and stuff like that and then we just compile and run this little C code file it's a single file of C at this point and uh I think it's like 2,000 lines or something like that if I remember correctly and you run that program and it like does a little training and outputs on Shakespeare at the end and then we can verify that pach code is identical to the C code and everything is great and we're just running in C and at this point I'm actually feeling quite great because this is amazing so if we have a single file C there's no dependencies whatsoever it compiles instantly runs instantly all the memory is just allocated in a single blob so if you start stepping there's no way you're going to Mo later it's all pre-planned it's cly deterministic it in principle can train gbt2 it's complete it will trange gbt2 you just have to wait a long time and it can run on a potato it can just run on anything just a single pile of C with no dependencies and in principle this could run this would be a great candidate to run on a on probe uh because in space if we just Harden it a little bit more uh because you know you're not going to ship High torch code on the V probe but I think LMC is a [Laughter] great so I was feeling great at this point uh on side note by the way all of this work I describ so far vacation while I was jetl in males so uh I basically it's perfect because you wake up at 1:00 a.m. and there's nothing to do uh so you write stuff like Ali and then in Sunrise you go do all the water activities so that is the Villa where most of el was uh was train uh so that was perfect this is a picture of it and this is a this is a I think the Moon is about to set and the sunrise is about to happen this is a recommended way to software develop okay so now we have C code but it's inefficient so we'd like to run it faster for that we reach for gpus so we need to convert all of our CCO to GPU so this is where we go the dev Cuda part of the Reco and we start to develop all the kernels so here's the layer forward pass as I mentioned and now we're going to develop a number of kernels that had the identical functionality but now run on a GPU and they're going to be faster and so usually we have versions 1 2 3 4 5 6 Etc and these are all different kernel implementations they're a bit faster usually over time but they match uh the specification exactly and give the exact same numbers uh so we develop all those layers and P them to Cuda and this is I don't know this is one of the kernels basically the point here is the first kernel is Trivial to do usually because you're paralyzing over batch and time and then you're basically copy pasting the C code into your Cuda Kel and you're already getting speed UPS because you're paralyzing Over The Bash time tokens and each thread just handles a single out element so the first Jal is usually trivial but then optimizations can be pretty elaborate so by the end we get to Kernel 6 for example in layer norm and we're doing a lot of things that are a bit more complicated so we have some uh you know um More Produce operations we have some uh we also communicate through shared memory through gber memory we're orchestrating it correctly uh cach screaming hints and a bunch of little tips and tricks for uh dealing with everything um and I'm going to go into a bit more detail later but this you can get arbitrarily complicated here writing the qua code one thing that I sort of found in in this project is that it's not exactly trivial to learn Cuda unfortunately and it was like a little bit harder than I expected um I I need some Cuda going in but getting better at it I think is not trivial I think some of these books unfortunately are a bit out of date as you might know uh pmpp is actually quite good uh but also I think still kind of like U mostly on the beginner level because a lot of the Cuda code that we ended up developing in the lifetime of the LMC project you would not find those things in in this book actually uh so a lot of the kernels that uh we ended up uh add would just not be covered and then on top of that you have this Cuda C++ programming guide which frankly is not exactly readable for someone who is like a bit new to that uh to the to Cuda uh and then you have this amazing block post from Simon Z anthropic that is like way better than anything we deserve just like randomly on the internet so that was incredible and if there was just more of that that would be so amazing but uh yeah so I think I found it a little bit difficult um but uh I mean I'm hoping that things like fud mode can can definitely speed up uh the uh availability of writing Cuda okay so next what happened is uh I was basically struggling with the Cuda code a little bit and I was reading through the book and I was implementing all these K kernels and they're like okay qu kernels but they're not great and so uh a team of Avengers assembled from the internet when they saw n start contributing so specifically uh Eric AR alexar like I would say core deps of L.C and contributed a ton of work to L.C and they they started to like really optimize and write all these kernels and this was incredible to watch and learn a lot fromont and there's many more Ros wheeler and Chinle and a few others uh there's over time we have 60 contributors to the Eli project shout out to Landa for sponsoring Eli uh they contribute compute so that we can run and optimize all these kernels uh so it was amazing for me that people just came in the internet and helped out on the project and you know this is one of the favorite things that can happen my favorite things that can happen with an open source MIT license dep people just come from the internet and help contribute it's amazing Okay so we've converted all the layers to Cuda we have now all the kernels and we can now train on a single GPU in F32 so far so that's great so from then on we start to make more and more optimizations so number one we don't want to have map m in fp32 when you roll your own code uh we actually switch to kuas uh step two we don't want to write our own flash attention I think that would be pretty complicated turns out CN has a very good flash attention implementation so we switch to that um next you want to definitely reach for mix Precision uh so that uh to speed up the code so you want to go over all your tensors for parameters and also for activations and so on and you have to start to think about okay which ones are in float 32 which ones are in FL 16 and what procision are they in and then we do all the conversions automatically so we uh reached for that and implemented that there's many many other optimizations that we ended up implementing over time so as an example we did all the kernel fusions uh different recompute settings to Rec compute a piece of the forward pass during the backward uh we um there's been a lot of optimizations from Eric especially on minimizing the amount of memory that you need during the macro pass um we have this like packed 128 data structure which basically in our experience forces the compiler to use the 128bit load and store instructions that are available but somehow the compiler is unwilling to use in many cases uh so I think Arun did a lot of uh work here where you just look at the SAS and you look at um the SAS the assembly and you are looking at what the structures are being used for your Loop and you figure out that okay this should be a 128 bit load to store but happens to be a 32 bit or something else because something in the end CC compiler is not going very well so it we found that this data structure kind of forces compilers hand a bit more um we implemented all kinds of K streams to overlap the part of the computation and this ended up creating like a total disaster um and so that's why I scratched it out because at one point of L.C as Arun would say I basically went in and I red it from orbit I just went in and control up for allions of stream and I just delete delete delete and basically I deleted all the streams made everything single threaded because we ended up getting all kinds of really weird race conditions and errors and so on I just didn't want to deal with it so L.C is not actually as overlapped as it could be but it's just like it's it's too much complexity for not enough gain at this point so uh but maybe the maybe we can slowly reintroduce some of it we have stochastic rounding we have full determinism full determinism turns out to be pretty hard because some of the kernels complexify a lot because you can't use atomics like the encoder backward was especially crazy because encoder backward um is Trivial with atomics but non-trivial without it anyway so a lot of the optimizations went into with with a lot of uh efficiency and determinism in Mar and accuracy like stochastic and so on next you want to use multiple gpus not just a single GPU so this is where you uh bring in nickel you start to um do alberu between all the different workers and this is where you also start to reach for started Optimizer State uh 01 or basically take your Optimizer states which are in float and these are really large um buffers for ADW and you can actually spread out a lot of the stuff across all the gpus and it really helps to keep uh your requirements down per GPU in terms of memory so very helpful to reach for that so currently Al that's uses 01 which is a shed Optimizer State there's a PR for 02 but I don't believe I merged that yet because it gets a little bit messy but might be merged eventually a lot of LMC just kind of like um balancing the Improvement and speed with the complexity of which you're actually introducing and so I've actually rejected a lot of PRS because of that because the code starts to get crazy and I think that decreases the amount of people can be onbo it uh and then after multip you have multinode so now you are running across multiple machines you have to make sure that you synchronize all of them that they can find each other and so on so we implemented all that and where that lead us to is that we can actually turn qp2 and we can actually reproduce it after all of that work so there's a post and discussions of LMC we can train the 1.6 billion gbt2 which was state-ofthe-art llm as of 2019 or so uh and you can train it on a single note of h100s in about 24 hours and that cost roughly $600 and the way you do that is it's extremely dependency free there's no need for python no need for pytorch uh so uh you do need CNN which is the most heavy dependency but CNN is optional so if you'd like to roll your own manual attention that is possible in elti but K is kind of like the haras dependency but after that it's just a bunch of C code you compile it and you run it there's no need for really anything uh so there's no need for cond environments pit installs there's just nothing it's just amazing uh and then you compile your code and you run it and it starts stepping you wait 24 hours and then this is it stepping print some Diagnostics we have almost 50% mfu here on one uh node which is quite good um and you get really nice plots and you beat gp2 on hella swag and basically this just indicates that the optimization went well no crazy numerical issues PES or anything like that for this size and uh yeah achieving a really good model in L in ly um we can still compare to pych because remember we have pych implementation for all the stuff in parallel on the side and so you can run the equivalent training room almost in p and we can compare the two implantations side by side and in particular at the time of writing that post and I don't know if this has changed because the p continues to optimize things over time but at the time of that post we were using Al that in 30% less memory and we were 20% faster in training just the throughput and I don't know if I fully super duper optimized the P implementation I did my personal best U but uh this is we were able to I think be P to in training of specifically gpk2 in Ling if you want to train anything else you're in a lot of trouble you have to change your code a lot and doing that and I'll come back to it but for gbd to training were better after all that work um and it also compiles and rs much faster which is beautiful toor compil takes like quite a bit of time like a minute or something you're just waiting uh so that's also something that I personally don't like to work with usually okay so looking back around turns out it wasn't all that simple there was a lot of stuff involved and it took a few months for a few people um but uh it was fun we learned a lot and we friends along the way this is the El C you know so it was it was great um ongoing work uh we are adding lava 3 support we actually thought maybe would have it done by today but uh there's a few more uh few more a little bit more work to do but we will have Lama 3.1 um training in very very soon uh we will have fp8 support So Arun has been working on this uh and and uh there's a big PR that's coming for FP support which is also interesting um and there's a lot of notable forks. they're allo uh the AMD Fork is very active as far as I understand and and quite good I think also the C++ uh could Fork is is quite nice and so a lot of folks um so encourage you to also Fork L.C is fairly readable I think I try to keep it clean well documented I think it's pretty well understood what's in there it's only maybe like I think 3,000 lines of code of basically C mostly uh and one more thought I think that I wanted to get across is um it wasn't all that halfhazard to start the project I had another uh motivation for starting the project and that's that I think um I what is like if P torch is especially torch compil a bit like DCC for software 2.0 it's a compiler then LM is a bit like writing assembly where doing everything manually right and basically I think um we wrote LMC as multiple people over a duration of 3 months and got something that was faster than pytorch in a specific setting of gbt2 train and so what this this exercise basically proves that this is possible now the problem is seem to spend multiple people several months but if Els are about to be become much better at coding over time then I think you can expect that the llm could actually do this for any custom application over time and so the llms could could act as a kind of compiler for any custom application you're interested in they're going to do all the llmc work and they're going to Output a binary uh that you can compile and run for your specific applications so I don't actually know if we like the use of python and Pyro and everything else is just a crutch because we humans are finite with finite knowledge intelligence and attention uh but actually don't you want to write all code in custom C cels and so on like maybe and so um the other thing that I think is interesting is the L might be useful because in the early stages of these llms and their intelligence uh they might not be able to write this code from scratch if you just prompted them write gbt2 andc you probably won't get lm. C but you're a lot more likely to get it if you put lm. C in the context of a s LM and you can expect that the few shot learning would be very helpful for the LM to basically give it example code and so I think ly could be very useful for this example code to give to the elm as they're about to write of our custom applications and so I think this is actually not unlikely to happen um yeah this is kind of likely to happen so I think software development in general will probably change a lot and to me that needs an exploration of whether this is even possible because if it is possible then maybe this this is what's going to happen so yeah that's it [Applause] thank all right like well the morning session talk