SBTB 2023: David Hall, Levanter: Legible, Scalable, Reproducible Foundation Models with JAX
[Music] hey everybody uh so yeah I'm I'm David Hull I'm the research engineering lead at Stanford center for research on Foundation models uh today I'm going to talk to you about Levant which is our framework for training uh uh Foundation models so I don't think I really need to tell anyone here that Foundation models are already having a massive impact in the world chat GPT reached 100 million users after just two months uh making it one of the I believe the fastest uh adopted product of all time Congress is obviously very uh and and the EU are very interested in uh regulating this space um and of course these models what one of the things that distinguishes these models is how much compute they require and therefore how much cost they require GPD 4 is rumored to cost around $100 million and we see similar size investments from companies like Google and anthropic and meta and you know pretty much everyone else these days um and that kind of raised a question of you know oh mik mic up great all right that raised a question of what should the role of Academia be in this new world uh we obviously can't train even even a university like Stanford can't train a100 million uh model on the regular um so you know how can we still be relevant in this world so in 2021 Stanford established the center for the research for research on Foundation models and coined the term Foundation model um and it's a massive interdisciplinary uh collaboration with uh you know across 10 more than 10 departments with over 30 faculty uh a few hundred students and postdocs um and uh they you know what we do at crfm is uh divide into three groups there are uh applications like work uh working in say biomedical domains where we have unique data sets uh we can collaborate with the Stanford hospital to develop models that will help their them improve patient care um there are of course deep technical advances things like flash attention uh came out of Stanford um and we can also um you know think through uh social responsibility what should how should we how should policy makers think about uh the role these models have in the world and this gives us some amount of uh and as by being a non-commercial entity we have some amount of objectivity that other that for-profit enterprises don't have okay so um crfm has an engineering organization it's very small it's just two of us uh plus a few Ras um research research assistants um and the you might wonder like why have a re an engineering organization at at a center like this and it it has to do with the scale the rate of innovation and the rate of investment um as it grows Academia has to adapt to keep up um and it requires a lot of engineering infrastructure and resources uh just to stay to build models and stay relevant uh we have two projects uh we have Helm which is a platform for uh evaluating large language models um uh and it's uh it sort of solves the MN problem so that we've we've evaluated I think it's over 60 models now on more than 40 benchmarks uh it's uh you know very cross cutting and holistic and we also have lant which is the focus of the rest of my talk um and it's a platform for training large language models uh and other Foundation models and its goal is to be legible scalable and reproducible um and one of the fun features of lter is actually as fast as many of the fastest commercial Frameworks that are out there including those from Mosaic ML and even those from Google okay so let you know what do I mean by legibility scalability and reproducibility for legibility uh we want uh to demystify what goes into training a large language model or any other Foundation model we want the code to be easy to follow um and we want to be able to separate the logic of the model from the logic of scaling the model uh both of these are very important features but they often get tightly coupled and it makes the C the code harder to follow and harder to adapt um we also of course care about scale um the goal we set for ourselves is being scalable at least up to about 20 billion parameters uh ideally more but this is the largest that we thought we had a chance of training um on the budgets that we have um and it should also be fast in the same ballpark as the best commercial players we originally started out with the goal of it being not embarrassing but um we actually did better than that um and then finally reproducibility um we want to be able to reproduce our results exactly uh we want want bitwise determinism meaning that you get the same results every single time um even in the face of preemption and restarts from checkpoints because most of our compute uh comes from a donated compute and is preemptable okay um so lavant is a new Jack based code base it works on GPU and TPU uh it uh places a high premium on legibility without compromising scalability uh and the performance as I've said is competitive with the largest libraries even those whose m goal is performance where they what what they do is brag about their performance and of course we have bitwise reproducibility so you get the same results every single time okay so legibility so this is um some fairly typical tensor code it happens to be written in Jacks uh but if you've written numpy code or P torch code this code should look fairly familiar um the goal is this is just evaluating the the loss the mean squared error uh between some uh Target y some features X X and some weight Vector W um and this code looks fairly straightforward and uh standard for anyone who's written this kind of code before um but it's actually wrong um and it won't throw an exception it will just silently give you the wrong answer um and it's it can be tricky to see why um and the answer has to do with these shapes of these tensors so Y is 128 wide uh but x * W is 128 by 1 which sounds the same but if you no numpy or Pi torch or anything else when you put together 128 and 128 by one uh two two ters of those shapes you end up with something that's 128 by 128 so you end up evaluating all pairs of um of uh differences and you end up with some meaningless number that doesn't isn't useful um and this kind of bug shows up over and over again in deep learning code or in any kind of numeric code um I've I've had this happen you know over a dozen times in my life I usually lose a day to it every time um and I I it's it's a it's an endemic problem um and also this code tends to be I think illegible so this is a library called MPT uh it was designed to be pedagogical um and I I even as someone who's worked in the space for quite a while now still have a lot of trouble understanding what's going on um there are comments but they're prone to bit rot um and you have to like really think hard about like what's being swapped here what does negative2 mean what do negative 1 um what are we normalizing over and then again what are we swapping and then we have these views um at the end that are very normal but they make it they hide bugs uh if you forgot that transpose your code would just silently generate the wrong answer just as we had before okay um so tensors with positional axes are bug prone and and less maintainable in the same way that Assembly Language is bug prone uh and less maintainable than higher level languages uh positional temper tensors typically overly prescribe the exact set of dimensions of tensors reducing the uh the reusability and modularity of tensor code um and so what we've done is we built a new library called Holo that's a name tensor library that improves readability reduces the risk of bugs and makes it easy to scale model training to hundreds of tpus or tens of billions or parameters so we're going to make things nicer to read and scale all at the same time okay so the the centerpiece of HX is this named axis where you pair a a name describing the semantics of some Dimension with the size um the sizes are the usual things you're used to in deep learning um and and the names are the things you would find in the papers but often don't make it into the the code um and we can create uh name uh uh name tensors using a very familiar syntax uh and we can even do operations like matrix multiplication uh the same way and if we take our example from before we can write code that looks basically the same um but we've defined way the ability to make that broadcasting bug um and and so we've eliminated an entire class of bugs and made our code easier to read all at once um and if we take a look back at this attention implementation where it was you know extremely hard to uh follow all the the indexing juggling we can instead have an implementation that almost looks like pseudo code um but it actually implements the exact same algorithm actually runs and ends up actually being faster than the code on the left okay so that's cool um legibility and expressivity are important goals uh but we still to be useful we have to be scalable and efficient as well uh but as I've said name tensors actually give us both um in fact it'll actually make scalability even easier to implement and to understand um so probably don't have time to go through the whole concept here but if we think about uh our our accelerators our tpus or gpus uh we can think of them as actually being a named array all on their own except instead of having uh numbers and then they have they have gpus or tpus and if we take a matrix uh like this one we can think about dividing it evenly across all of the tpus uh along some axis so this is dividing up the uh our input embeds diens uh input embeds tensor along the batch axis um and this blue uh yeah and uh we can take that conceptual I you that that idea and we can implement it in code in this exact way by um saying we're going to to map the batch axis to the device axis of our devices and we're going to Shard our uh data across this mapping and just like that we've made our codes uh data parallel and it will it can run on as many tpus or as many gpus as you as you can get your hands on that's pretty nice um but we can actually do better so um one of the the state-of-the-art method for training extremely large models these days is called something called fully shed data parallel you also may have heard of it called zero um and um it and uh the Insight it has is that typically in normal data parallel training you have a full copy of your parameters on every single TPU um and you split up you only split up your data across uh each TPU uh with fully shed data with fully sharted data parallelism however you instead split your parameters across some axis here the embedding Dimension uh and you end up with a just a single slice of every parameter on every TPU and whenever you need to do compute you do a all gather a network sort of communication so that you send all the relevant parameters to all the machines that need them and do your compute uh this is very complex uh operation that if you look at the pytorch implementation there's you know th tens of thousands or uh lines of code to get this to work um and we can add it uh to any modeling code in about 10 lines of code in h x now we just add a parameter mapping and uh we uh Shard our model just like we sharded our uh data and um we can also Shard our Optimizer States and uh in addition the data from before and now we just um uh give a little bit of extra help for uh how we do the computation we want to Shard the data during computes to look like the data uh charting and uh just like that we've implemented fully charted data parallel okay um so I've shown that you can scale you can have a legible implementation and you can still design for scalability uh but scale means more than just you can referent fstp it actually has to be fast and it actually has to scale um it turns out that that it does so this is something called Model flop utilization uh don't worry too much about what it means but basically means like how good of a job are you do are you doing using the hardware you have available to you um this is measured on a V3 256 tpod um which is like 50 a100 worth of compute um and if we look at some comparisons um it's not super scientific because they're different Hardware but we're achieving mfus that are in the same ballpark and and sometimes better than uh state-of-the-art um uh implementations these super large models tend to have much l lower uh utilization so the comparisons of palm and Megatron aren't super Fair um MPT that pedagogical implementation from before um is actually hovering around a little under 20% so we're we're actually doing quite well um uh by this metric okay um I've got five minutes left I think um so I can talk about reproducibility but short answer is pytorch code is actually very hard to get reproducible this is four runs that were launched on the same machine with the same random seed with all the determinism flag set and it's still not deterministic um and you know you may think whatever but it turns out that it actually can really matter especially when you're training these really large models they tend to be very unstable when you're training them um and this is two runs that were started from the same random seed uh in pytorch and one of them uh blew up and the other didn't and we have no way of tracking down exactly what went wrong because because we can't reproduce the run you have oh 15 minutes oh great doing great all right um cool I I'll have time to go through fun examples then um so [Music] um let me uh so uh great so in Levant however we get bitwise reproducibility so this is again four different runs same random seed uh but this time we're actually going to use different config we're going to use different tpus uh different TPU machines same configuration um and we're actually going to kill the jobs and restart the jobs and we end up with perfect reproducibility even in the face of these restarts these are bitwise identical results um and so that allows us to to do to do science okay um so uh let's talk about some of the cool things we've been doing at crfm with uh with lant uh so the first thing I want to talk about is a new Optimizer called Sophia um so most of you who' have done deep learning have certainly heard of Adam it's been the State ofthe art deao Optimizer for the last nine or 10 years now um and um a student named Hong and his adviser tangum ma um developed this new Optimizer called uh called Sofia that's a new second order Optimizer that uses an approximation of the hesen which is to say the second derivative uh to uh improve the optimization uh to proove uh the optimizer sort of estimate of what direction to follow um and you end up with an Optimizer that is two times faster than atom uh and they started out by training small models of about 100 million parameters to maybe 700 million parameters um and uh they couldn't go any higher and when they put out the results they got a lot of feedback from industry basically saying we don't believe that this will scale this is not something that's relevant to us we see this kind of stuff all the time so instead they took Levant and um ran uh uh implemented Sophia in in about a day um and scaled it up to one and a half and then 6.7 billion parameters and showed that it actually does scale and that even out to um uh 100,000 steps or I forget exactly how many tokens but uh probably 50 billion tokens um it's two times faster uh than atom so we have this new Optimizer we developed that's easy to implement doesn't require additional memory and ends up being two times faster uh in terms of Step count and about uh 1.6 or 1.7 times in terms of wall time something like that uh we also have some fun um uh uh we also implemented controllable music generation so John Thon is a postto here at Stanford um and he developed an approach to um an approach to music generation that allows you to give a prompt to the model and it unlike with a with an L an llm The Prompt can be not just a prefix but it can actually be uh any infix or it can be just a single voice a single instrument from the uh that that a human made um and then the model can generate the rest and can generate a rich accompanyment with multiple different instruments all at once uh I'm guessing sound probably isn't going to come through so I'll skip the audio demo um but the the upshot is that the results are um uh are uh for company generation are uh every bit as good as a or rather are indistinguishable from humans at least According to some carefully vetted uh turkers um so uh that's pretty cool um uh and then uh just to talk just a little bit about some of the uh uh the new things we've implemented uh so uh We've improved the documentation things are um are should be easy to get started out of the box on GPU or on TPU um we also uh have released support for llama and Laura so you can trade in alpaca um with just a couple lines of code um and we have an implementation of flash attention two that allows you to scale up sequence lengths from the normal like 20 48 or 4096 that people do these days all the way to 32,000 um without um you know without too much trouble um and uh coming soon uh we're working with Nvidia right now to fix uh support across multiple gpus uh or sorry across multiple machines uh currently you can uh run Levant and it works great on a single machine um with you know eight or even 16 gpus in it uh but where there's some Network bugs um we're also working on uh adding uh Dori which is a new technique that was developed here um at Stanford uh that will automatically tune the mix the training mixture to optim to be optimal for uh the kind of model you're trying to train we're also uh improving the flash detention implementation so that it can automatically tune itself um for whatever the hardware uh is that you're using uh and soon we'll be adding Vision Transformers so that we can train uh not just Lang large language Foundation models but uh Vision models or hybrid language Vision models uh since I've got a little bit more time I'm going to let's get backpacks I'm going to talk about Lumia so one of the projects we're training uh at at Stanford and collaboration with Stanford hospital is Lumia and it's uh it's uh we're going to take about 3 terab of patient data that's been carefully anonymized um that includes all of the notes and test results and also all the diagnostic codes that uh doctors use um and we're going to use it to uh improve diagn diagnosis do what's called risk gratification trying to figure out who's likely for a repeat visit to the hospital um do things uh like history retrieval so that doctors don't have to go scrounging around in the medical records for forever um and uh and we're uh started kick we started the training run for uh sort of the first phase of this um just the last few couple days ago I'm really looking forward to the results of this uh in the next few months um yeah I I'll I'll cut it here um yeah thanks we've got some time for questions I can run mics so what do you think that is the Innovation behind of the framework that you guys building now that other libraries like caras are multi backing and you can pretty easily use a friendly API to create models and just switch from one back to another you know like fighter or jacks or what do you think that this could be the benefits of use one a day oh so so it was a little hard for me to hear just with the the the the um the speaker um I I think the question was what was the motivation for choosing for choosing ja and choosing our own framework um so this yeah it's a great question um you know the reality is that most people in Academia are using pytorch still um but uh Google uh generously gave us access to um a large number of tpus uh for free and we started off trying out pytorch on TPU and um support was preliminary I guess maybe the right way to put it um and it just it didn't seem like you could train anything of any size um using it um at the time um and so Jax allows us to be opportunistic with our compute uh you know Nvidia as I've said is working with us to to make mon sort of first class uh for their on their hardware and you know we're you know possibly going to talk to other organizations like Intel and arm who are building their own accelerators and it allows us to sort of take compute wherever it comes from cool thank [Music] you