Intro

In this blog post I will be writing about some of my high level learnings from a recent LLM conference I took part in. These are some notes for myself so I don't forget. These concepts are basic and high level but are details I was missing before when thinking about fine-tuning LLMs. Most of the details here were inspired by one of the conference speakers: Jonathan Whitaker's ---> talk.

Bits and Byte

You have probably heard the saying, "In the computer it's all 0's and 1's". Well, those 0's and 1's are called bits. A bit is the smallest unit of storage and simply stores a 0 or a 1. A byte is a group of 8 bits together.

• 1 byte = 8 bits.

All storage is measured in bytes.

Memory Usage During Inference

I'm running this code from a local Jupyter notebook in Pycharm connected to a A6000 running on JarvisLabs. I used the axolotl template.

Yes, that's what we thought.

Let's run some inference with the model.

Now let's try loading a quantized model using BitsAndBytesConfig to a load a model in 4bit precision. Note that the model weights are stored in 4bit precision but the computations are done in a higher precision. Here we are specifying that the dtype for computations is bf16. During inference, as well as training, the weights of the model are constantly being dequantized (from 4bit to bf16). This can be done for specific layers at a time during forward and backward passes, to keep memory requirements low. The computation happens in a higher precision.

Here we specify that the model should be loaded in 4bit and the computations done in bf16. Since we are using 4bit we expect each parameter to use 4/8=0.5 bytes so we should be using less than a GB of memory.

You can experiment with different types of quantization.
In this next example we load the model using:

• NF4 quantization.
• bnb_4bit_use_double_quant which uses a second quantization after the first one.
• bfloat16 for computation

I really like the high level explanation of quantization from this post by Patrick von Platen. In general, when running inference with quantized models the steps are:

• Quantize all the weights of the model and load it (for example 4bit).
• Pass through the input sequence in bf16.
• Dynamically dequantize the weights to bf16 layer by layer during the forward pass
• Quantize the weights back to 4bit after the computation

So if we want to do $Y = X W$ where $W$ and $X$ are the weights and input sequence respectively, then for each matrix multiplication we do:

$Y = X \cdot \text{dequantize}(W)$ ; $\text{quantize}(W);$

For this reason, inference is usually not faster when using quantized models. It's slower. It is good to remember that quantization is a tradeoff between memory usage and output quality, as well as possibly inference time.

Memory Usage During Training

This is not something I know a lot about at the moment, so I can't go into too much detail. But I can cover some high level basics that I learned recently. There are three main areas which contribute to the memory during training:

• Model Parameters
• Gradients
• Optimizer State

There are other things to consider such as the activations which tend to dominate the memory at larger batch sizes and context lengths. But let's ignore this to keep this simple and high level for now.

Suppose we want to fine-tune llama3-8B in bfloat16 with the basic Adam optimizer.

• Model Parameters: 2 bytes per parameter for 8B parameters is 16GB
• Gradients: To store the gradients for each tunable parameter is 16GB
• Optimizer State: needs 2X the size of the model, to store first/second moments, which is 32GB.

So you would need at least 64GB to fully fine-tune llama3-8B in bfloat16. What can we do to fine-tune the model with much less memory?

First, we can quantize the model to 4bits. Then llama3-8B would take up 4GB of memory for the model parameters. That is a 4X reduction! But when training we don't quantize the trainable parameters or gradients because the training would not converge. Training still needs to happen higher precision. This is where PEFT methods come in handy, such as LORA and QLORA. Let's consider QLORA since we are discussing quantization.

With QLORA we add a set of trainable adapters whose parameters take up a much smaller percentage of the total model parameters. We can freeze the entire quantized model and keep it in 4bit. We can store the corresponding gradients and optimizer state in higher precision. This is possible because we are only dealing with a very small percentage of the total model parameters that are trainable.

In my last blog post I fine-tuned llama3-8B using the axolotl Library. It was configured to use QLORA with the model parameters in 4bit precision. It was using around 15GB of memory during training. There were some spikes due to me loading a model in a different python session, so just ignore those.

Inference with Axolotl Fine-Tuned Model

This was the corresponding bits and bytes config for my Axolotl fine-tune. As said previously, it was trained using QLORA in 4bit precision.