Efficient Fine-Tuning for Llama-v2-7b on a Single GPU

Sam.

Hi everyone, my name is Diana Chen Morgan and I'm part of the DeepLearning AI team bringing you all together for all things AI community and events. Today we are very lucky to have a workshop with some special speakers from Predibase. In this hands on workshop, we'll discuss the unique challenges in fine tuning LLMs and show how you can tackle these challenges with open source tools through a demo. By the end of this session, all attendees will understand how to fine tune LLMs like llama 2 on a single GPU. Techniques like parameter efficient tuning and quantization and how they can help and how to deploy tuned models like Llama 2 to production. With continued training with RLHF and how to use RAG to do question answering with trained LLMs. This workshop will be recorded and remain live on our YouTube channel and you'll be able to access the notebook in the description of the video and the chat. Everything covered in the workshop is presented as continued education from our existing AI short courses and we'll be dropping the link to access all of them in the chat. To start, I want to introduce our first speaker, Piero Molino. Piero is the co founder and CEO of Predibase. He was one of the founding members of Uber AI Labs where he worked on several deployed ML systems including an NLP model for customer support and the Uber Eats recommender system with graph learning and collision detection. Later he became a staff research scientist at Stanford University working on machine learning systems. He's the author of Ludwig AI with 8,900 stars on GitHub, an open source declarative deep learning framework. In 2021 he co founded Predibase, the low code declarative machine learning platform built on top of Ludwig. Hey Piero.

Hey Diana. Thank you so much for having us.

Yeah, we're so excited to have you here. And for our second speaker we have Travis Sedair. Travis is a co founder and CTO of Predibase, a low code platform for predictive and generative AI within the Linux Foundation. He serves as lead maintainer for the Horvod distributed deep learning framework and as a co maintainer of the Ludwig declarative deep learning framework. In the past he led Uber's deep training learning team Deep learning training team as part of the Michelangelo machine learning platform. We are so excited to have them here today and we can't wait to see what they have in store for us. So Piero, take it away.

Thank you so much Diana. I'm also super excited and I've been part of this community from the other side of the deep learning AI. So contributing to it and somehow it's really great for me. So let's start talking about efficient fine tuning for llama 270 billion on a single GPU because it's not that easy to make it run on a single gpu. Diane already introduced us, so I will skip this slide, but let's talk about what we'll cover. So first of all I want to give you like a little bit of an overview of how I think about fine tuning and why would you want to fine tune an LLM to begin with, I will introduce Ludwig, which is the low code framework for building custom AI models that I developed when I was at Uber. And now it's like the foundation of our technology at Predibase. And also what are the challenges of fine tuning in particular the memory bottleneck and how to overcome them through different techniques to squeeze the 70 billion parameters into something that fits into the 16 gigabytes of VRAM of commodity GPUs and hardware. And so we will cover half precision quantized training, low rank adaptation called Lora and Qlora, the quantized version of it. And we'll also show you a little demo of how to make this work on your own. So why to fine tune your own LLM to begin with? This is like a graphic that I created that basically shows a distribution of AI tasks within organizations. And with respect to the availability of data within organizations, on the left hand side you can see that if you have a lot, a lot of data in your organization, then you can even go as far as training your models from scratch. It will be pretty expensive, but it's certainly doable. But if you have in a reasonable range of data and already have pre trained models that you can use, you can fine tune them on your tasks to get really, really good performance. If you don't have a lot of data, the best choice that you have right now is to do in context learning through for instance retrieval augmented generation using general models. But what we will focusing on is the fine tuning part in particular because it has a really great trade off in terms of accuracy versus performance and speed with respect to the bigger, more general models in particular, because you can take a smaller model like llama2.7 billion and fine tune it on your data to make it perform as if it was a 70 billion model or even better than that. So first of all, let me tell you a little bit about Ludwig. So that is the tool that we're going to use for showcasing These capabilities. My experience when I was was at Uber working on different machine learning applications was that through all of them, like for instance the intent classification model for customer Support, the Fraud Prediction 1 and the Product or commander for Uber Eats that Diana mentioned. In all of these cases there was a lot of code that needed to be written and the development process took a long time, not really for writing the code, but for iterating and experimenting. And then the deployment of these models also took a long time. So I thought there could be a better way for doing this. That's why I came up with Ludwig, which is again an open source declarative machine learning framework, in particular tailored for deep learning use cases. And what does it mean to be like a declarative machine learning framework? It means that you can specify your deep learning pipelines by just using a configuration file. And the configuration file only contains information, only needs to contain information about the schema of your data. What are the inputs, what are the outputs and what are the data types associated with them. In this example you have one input sentence with type text and one output intent with type category. And this six line configuration matches 100% of the performance of the intent classification model. That was the one that I developed when I was at Uber. It's very easy to iterate over these models and get started building models in this way. And you don't need to write low level machine learning code. At the same time though, you have all the expert level of control that someone who has been studying the deep learning AI classes gains. Right? You can change every parameter of your models from the training parameters to the architectures that are used, all the way down to the single hyperparameter of the architecture, like the single activation of the single layer and pre processing parameters. There's more than 1,000 different parameters that you can change and modify. And this makes it very easy to iterate of your models because you can just change one line in a configuration file instead of changing multiple lines in your code. Finally it's extensible. You can write your own Python classes, give them names, and if you use those names in the configuration then you can extend the system as much as you like. Finally contains also advanced functionalities like hyperparameter optimization, set of the app models that are already baked in and also distributed training all out of the box. So with this kind of approach now new people that before were not empowered to build deep learning models, but just writing 15 lines of a YAML configuration file in one day, they were able to train a model for website, personalization so these were engineers at Uber, not product engineers at Uber, not deep learning experts. And in one week they managed to deploy it into production. And so this kind of approach makes it substantially easier and faster to develop deep learning models. The reason for that is that at the core of it, there's this architecture which makes it possible to have multiple inputs and multiple outputs of different data types, like for instance, categories numerical and binary values, or text, image and audio, and to produce as output text numerical binary values and so on. The inputs are first pre processed, then encoded with a piece of a deep learning model that produces vector representations that are then combined by one single component to be provided then to the decoding components that produce from these vectors the final predictions of the different types. And you can imagine how mixing and matching different types of inputs and outputs can lead to different machine learning applications. So for instance, if you have category, numerical and binary values as input and numerical values as output, you're basically training a regression model. If you have text as input and category as output, then you're training a text classifier or an image as input and text as output, then you have an image captioning system and you can imagine how many combinations of types of inputs and types of outputs lead to different applications and how flexible these are approaches. Notably also if you have in particular text as input and text as output, you can definitely use an LLM for attacking these kind of tasks. And in particular, if you have this kind of data, then you can use it for fine tuning those LLMs. Here I'm showing you a configuration where you just specify the model type LLM, the base model which is La Manchu 70 billion 7 billion sorry, and inputs and outputs text features. And you can specify also the parameters of the trainer, like the learning rate, the batch size, the gradient accumulations, the number of epochs, and a few others like decay or any other kind of optimizer related parameter. Once you write a configuration like this, you can just provide it to the constructor of the lutbigmodel object and then call train, providing also the data frame containing your data, this will fine tune the entire model, both the embedding part and then also the prediction part. But if you wanted to fine tune a model to perform a new task, so you're basically replacing the head of the model with task specific head. What you can do is you can change the output features to contain in this case a category output, for instance a sentiment category output. And then you can specify the large language model as an encoder for your. For Your for your task, the rest of the training and the how do you create a model, how you train it is exactly the same. And you could do the same thing if you want to actually just freeze the weights of the large language model. If you do that, the advantage is that you can also cache the embeddings, which makes training substantially faster at the expense usually of some performance. Because in this case you are training just the classifier head and not the core trunk of the model. In all of these cases though, the problem is that these large language models are in particular like llama7 billion, are larger than what fits in the memory of a GPU. There are like newer and better GPUs like the H100 and the A100 that have quite, quite large VRAMs of 80 gigabytes. The problem is that they are pretty expensive from cloud providers. And also there's currently a pretty bad shortage of these GPUs, meaning that if you request from AWS a GPU like that, it may require quite some time to obtain one, if you will get it eventually. Right. And so the better idea is to use actually commodity GPUs like T Force or the RTX 4080, which is like a consumer grade one. The you could go out and buy. And the problem with these GPUs is that they have only 16 gigabytes of VRAM. Because of that, you need to be a little bit smarter in what you do in order to fit the models into the 16 gigabytes of RAM. And this is like what now Travis will talk about in his section of the presentation. And also we show you how to do it within Ludwig.

All right, thank you very much, Piero. All right, hopefully everyone can see my screen. Okay, so as Piero said, the main issue that we want to talk about today is how to make the most your use of very finite resources, which when we're talking about training large language models, is almost always the VRAM or the amount of memory that your GPU card has available to it. And that's in particular because your T4 or your 4080 typically has about 16 gigabytes of RAM. And for the purposes of this visualization, every one of these squares you can think of as being a gigabyte. But if you want to train, even a fairly modestly sized large language model like llama2.7 billion 7 billion parameters, when you actually want to go and train, ends up looking like something like this in terms of the memory requirements. And so if we break this down into three different buckets the model parameters, the gradients that are produced during training, and then the optimizer state that's needed to, to keep track of the training variables. Then what you find is that you need about 28 gigabytes just for the model parameters themselves, and then an equivalent amount for the gradients during training. And then usually you need about 2x the number the memory used for the model parameters for keeping track of the optimizer state. So in total, you're looking at about 96 gigabytes more memory than you actually have available. And so the idea of trying to pack all this into 16 gigabytes on the surface might seem like an impossible task, but that's exactly what we're going to do in this talk, is kind of break this down piece by piece until we get into something that is very manageable and fits on this single, single card. So let's start with the first piece, which is just loading the model in the first place. Right? So the very first thing you run into is that because you have 7 billion parameters and every one of these parameters is by default represented as a 32 bit floating point. 32 bit floating point means 4 bytes, right? So that means you're looking at 7 times 4, which is 28 gigabytes of RAM. So that's the first hurdle we need to overcome here is how do we pack the model parameters into a smaller form factor. So you may be familiar with the fact that there are different ways to represent a floating point. You know, floating point 32 is the one that we most commonly use because it gives you a really nice mix of the dynamic range, which is expressed as the exponent here in orange, as well as the precision, which is expressed as the significand, also sometimes called the mantissa here in the blue. Right. And so very natural thing you might say is, well, what if instead of using float 32 we go down to float 16? There's a trade off there, which is that if you just use float 16 by itself, you can obviously reduce your memory requirements by half, but you're going to non linearly decrease the range of values that you can represent because you're reducing this exponent that can lead to a very significant drop in how big of a range of values you can represent, particularly for your gradients, which can lead to things like vanishing gradients or exploding gradients, which can lead to nands and things like that. So one alternative that exists for people that have the hardware is to use a data type called bfloat, stands for brain float. Came out of Google is now officially supported on newer generation cards like the Ampere A1 hundreds and things like that. But for those of us still using T4s and lower generation hardware, you have to stick to something like float 16. So there's a trade off there. But if you were to go with float 16 or bfloat 16, that right there allows you to cut the memory down in half, 14 gigabytes. Suddenly you can fit everything on your single card and no problems. But very quickly you run into the next problem, which is the gradient. So, gradients. I'm sure most folks here are familiar with the normal stochastic gradient descent training process, but essentially what happens is during the training process, you produce a forward pass over the model. Those are what we call the activations. And then when you compute those activations, you want to then use the results to compare against the target that you're trying to predict. And then the difference between the prediction and the actual target is then used to predict produce these gradients, which are used to then do the backpropagation and then ultimately update the model. So the problem is that the gradients typically use the same data type as the original model parameters, and there's typically one gradient for every model parameter that you have, particularly the ones that you want to train and update. And so as a result, anytime you're doing normal gradient descent training, you're looking at about a 2x, just without any optimizer on the amount of memory required for training. So how do we go beyond this? Well, one approach is to say, well, is there something lower than float 16 that we can do, maybe 8 bit? And at this point, oftentimes a common technique would be to use quantization. So when we say quantization, what we mean is taking a continuous space, like the space represented by a floating point value, and discretizing it into a specific finite number of bins, which can be represented just as integers, kind of like an index, like 1 equals bin 1 and 3 equals bin 2, et cetera. And so naively, you could do this just by very simply slicing up the space uniformly and then bucking them all together. But in practice, you often find that values are concentrated in certain parts of the range. So for example, we have a higher concentration of values here towards the mean. And so if you were to just naively bucket things, you would lose a lot of precision towards the mean at the expense of, you know, in exchange for basically getting more precision towards the outer edges of the distribution, which is not really what you want. Right? You want to retain that nuance as much as possible. And so usually what you end up doing is calculating some statistics about what each bin represents in terms of what values are in there. And that allows you to then reconstruct this distribution when you want to go back to a floating point value in the future without losing much of the data. Right. And so outlier clipping is another approach that's commonly used under the hood to avoid this problem of completely shifting your buckets too much based on very extreme values. And so this seems like a very complicated process. But the nice thing about declarative frameworks like Ludwig is that we make it very simple. You just say quantization bits equals eight and boom, you get all this compression for free, essentially. So let's say you do that. So once again, we cut everything in half. So essentially, now we have only 7 gigabytes required to allocate the model parameters and an additional seven gigabytes required for the gradients. That means we're only using 14 of our 16 gigabytes. Right. So great. However, the real kicker is the optimizer. And I think a lot of people don't realize this initially, but if you're using a very common state of the art optimizer like Atom, Atom is great. People love it. It has become like a de facto standard in the industry. But it does have quite a high memory footprint, which a lot of people initially think about is very real. And in particular, as I said, you have two x the number of parameters due to the additional vectors that you're keeping track of for the optimizer. And that leads to very significant problems, as you can see here, where essentially all the parameters of the optimizer are completely out of memory. Now, one question you might have is why does Adam use so much memory? And this is the normal update formula for how you update the the weights based on the gradients. So these parameters here, the gradients, and then, you know, W's are the weights. And then these parameters here, M and V is momentum and variance. And these are two atom specific parameters that are essentially vectors the same length as the gradients. Right. And so we don't need to really get into the math here, but the important part is that this, you know, M momentum vector and this V variance vector, they're each 7 billion parameters in size. And so that means that that's ultimately where your 2x value is coming from in terms of 2x. The model parameters are needed just for the optimizer. So what can we do about this? Well, this is where a technique that you might have heard of called low rank adaptation comes in. And the interesting thing about low rank adaptation, or Lora, this was a technique pioneered by folks at Microsoft Research. I linked the paper down there in the bottom. Right. The interesting thing about this technique is that we often talk about it as a way to reduce the number of trainable parameters, to reduce the size of the footprint of the model weights when you save it and speed up training. And those are all true. But I actually think the most valuable initial thing about Lora is that it significantly reduces the number of parameters that need to be tracked by the optimizer as well as the gradients. And that's what ultimately leads to a very significant memory reduction during training. And just as a quick kind of refresher or kind of explanation of what Lora is really doing. So the key idea behind Lora is that when you're fine tuning a very powerful model like Llama, you don't actually need to fine tune literally every single parameter. There are typically some parameters, some layers, let's say, that are more important than others. Usually things like the parts that do attention and determining which tokens in the sequence are related to other tokens and what way are they related. And so the idea behind Lora is essentially to take those parameters, those matrices for the queries and keys and things like that, and then inject another kind of lower rank matrix beside it that initially acts as somewhat of an identity function. So it's not really doing anything. But then over time, when you propagate these gradients through it and update the parameters, the only thing you're actually modifying is this ancillary matrix that you've added side by side. Really the only parameters you need to update are here. And that means that the total amount of stuff space you need to save is just this little extra matrix. And that additionally means that your optimizer and gradients are only taken with respect to these parameters and not these big pre trained weights. The pre trained weights are essentially frozen. And again, this is all done in Ludwig by the declarative config, which is a single parameter, which is nice. So let's say we do this, let's say we use Laura and we now have added these additional low rank matrices alongside our existing parameters. What does that do to our overall memory footprint? And so it's not 100% like as easy as saying that the number of parameters is in the optimizer is 2x the parameters in the mall or anything like that, because the number of parameters that you add during the lower process can vary depending on some of the hyperparameters you set for it, like the R, which is how many the rank of the matrices you're going to add. And then you can also change which layers you want to add these low rank adaptation matrices to as well. So it's something that can vary a bit, but in practice it can be as low as 0.1% of all the parameters in the model, which is huge. And even in a more aggressive case, you're typically looking at something like less than 10% of the parameters. Usually like 1% of the parameters are actually going to be modified with Lora. And so one way that we might kind of split the difference a bit is say, let's assume that 1 gigabyte of memory is going to be used for these Lora parameters. Now, importantly, in the formulation of doing quantized training, normally what you do, and this was in the QLORA paper that we'll talk about, is you quantize the original model parameters, the ones in orange here. So these are represented in int8, but the LoRa parameters, those you typically do in FP16 or FP32, here we're doing FP16. And that's because since these are the parameters that you're training, you typically want a little bit higher precision on those. And that just. And since there's so few of them, it doesn't make that much of a difference anyway in practice in terms of the memory overhead. So 16 bit LoRa parameters here, plus the gradients for those 16 bit parameters. So 2x plus the optimizer state. And the optimizer state here, I'll say is usually FP32, which is why it's not just 2x the number of parameters, but 4x. So we have 2x the number of parameters times 2 when we go from FP16 to FP32. Back of the envelope math, you might be looking at about 4 gigabytes of optimizer state just using lora and int8. But there's still one more problem, which is the activations themselves. Activations are when you're doing the initial forward pass, effectively the memory overhead from activations. What comes out of each layer is going to be the size of the largest layer in your network, also times the batch size. So how many examples you're updating at once. Right. And again, this is not a kind of hard and fast thing because, you know, if you increase your batch size, you're going to have more activation memory to deal with. If you change some layer parameters, you might have more activations to deal with. But you know, back to the envelope math, you might be looking at about like 4 to 5 gigabytes of memory from activation, something like that. And so that is potentially another big thing that's going to cause you to have to deal with out of memory errors, even if you get everything else under control. So what can we do? How about one more trick which is can we go from 8 bit quantization all the way down to 4 bit? And so this was one of the most, I think, interesting and kind of groundbreaking things that came out of the Qlora paper. And so credit to Tim Detmers, linked his tweet there for the original image source here. But definitely check out the paper that he and others did on Qlora, which I think has been a really big game changer for this efficient fine tuning approach. So the way that this technique essentially works is that you go from full fine tuning to Lora, which only adjusts these adapter weights. And then with Qlora, what you essentially do is the base model parameters are represented now with just four bits. And that's because again, you're only doing kind of the forward pass on these. You're not having to update these parameters or anything like that. The adapters still in FP16, still 16 bit, but then now you have this optimizer state as well, which is FP32. And one additional thing that they did in the Qlore paper, which we won't get too much into here, but is one additional thing they did to reduce memory pressure, was introduce this paged atom implementation which allowed them to essentially offload the optimizer parameters to the CPU when needed to reduce the effects of memory spikes during training and effectively get this whole thing down to some very, very small form factor that can train on very commodity hardware.

Right.

But specifically we'll focus on the four bit part here for the purposes of the diagram here. And so, you know, the way they did this, they introduced a new data type. They called it normal float 4. Definitely recommend checking out the paper for more details on it. Effectively what they did was they defined like kind of a hard coded bucketing system. The main idea was to try to avoid the problem of not having a way to represent zeros. They kind of worked around that problem pretty cleverly. But essentially you do all this now, all of a sudden you go from int8 down to normal flow. 4,7 billion parameters can fit in just 3.5 gigabytes of vram. The lora params. And the gradients are still the same in FP16, your optimizer state is still the same here in green. And then finally you have, and apologies that this is five instead of four green squares. And then finally you have your activations in red here, which are just now at about 3 gigabytes. Because again, when you reduce the model parameters, you're also reducing the size of the activations themselves, which is a nice benefit. And so, all in, you might be looking at only around 13.5 gigabytes peak memory usage required here. So let's talk about one last thing here, which is your batch size. So typically during training, we try to pack multiple sequences into a single batch. One of the reasons to do this is less to do with efficiency of training, because often when you're training on small hardware, you know, you're lucky to be able to get one sample in a batch at a time, let alone two or four, etc. But one downside to doing that is it kind of, it can increase the variance in the training process in a way that's not desirable. So if you think about back to, you know, CS229 like, or you know, like kind of basics of training processes, you know, you have on one end of the extreme pure stochastic gradient descent. So just like one sample is one update all the way to full batch gradient descent. And one of the reasons why we don't like to do full stochastic gradient descent is it's a very noisy walk through the optimization space. You can oftentimes take lots of small turns in different directions, whereas as you increase the batch size, you take bigger steps and it smooths out the process. This sitting in a nice middle ground between taking these big smooth steps and taking these very jerky small steps. It's one of the arts of deep learning. But in general, you do want to kind of hit a sweet spot there in the middle, which is why we often use batch sizes in the area of like 32, 64, 128. But the problem is if you are only able to pack a single example into your model at a time for training, you end up doing this very janky stochastic gradient descent thing. So how do we work around that? This is the last technique I want to talk about today, which is gradient accumulation. The key idea behind gradient accumulation is that you can get the effects of training with a larger batch without the additional memory overhead. The way it essentially works is you do a normal forward pass. Here's a batch size of two in this case. And you can ignore the fact that this is a multi GPU example, we can just generalize this to one gpu. You do a backward pass, but then instead of updating the model params, you store it in a buffer and then do another forward pass and another backward pass and then basically sum those two together into a single vector, which you then if you're doing distributed training, you would then sum it once again in the Allredo step.