Transcription
Hi everyone. In today's video, I'm going to talk about how to compress and optimize your model. Let's say that you have trained or fine-tuned the model and now you're ready to deploy it for everybody to use, but you find that the latency is too slow and you want to make your model faster. Well, today I will show you four ways of making your model faster: quantization, pruning, knowledge distillation, and engineering optimizations. My name is B. I am a machine learning engineer and a PhD in natural language processing. Without further ado, let's get started.
With quantization, neural networks are large and take up a lot of space because they have millions or billions of parameters. By default, when you train your neural network, usually the parameters are stored in FP32, which means that each parameter takes up 32 bits. The idea of quantization is to reduce the precision of the parameters into a format that takes up less space, for example, 16-bit floating point or INT8. If we store all of the parameters in INT8 format, that means all of them are represented by integers between 0 and 255. Then that means we will save four times as much space compared to the original network in FP32 format.
There are several ways you can turn floating points into integers, and the most common way is called zero-point quantization. I'm going to step through an example of how this works. The reason this is called zero-point quantization is all the zeros in the original matrix are mapped to zero in the quantized version. We will see later why this is useful for sparse neural networks. Next, we take the maximum absolute value element and map it either to 128 or 127. In this case, the maximum absolute value element is -51.5, so this gets mapped to -18. The quantization has to be a linear transformation, so with two elements determined, the rest of the elements are determined as well. Finally, to get the INT8 representation, we add 28 to each element, so all of the elements are positive.
Two ways of doing quantization are quantizing the weights and quantizing the activations. In weight quantization, we store all of the weights of the neural network in INT8 format, and we dequantize the weights into FP32 when we run it, so that all of the data remains in FP32 format throughout the network. Since everything has been done in FP32, it is not going to be faster than the original model, but this is still useful because it saves space. For example, in mobile devices, we're making the model four times smaller is a significant improvement.
On the other hand, in activation quantization, we convert all of the inputs into INT8, and all of the computations are also performed in INT8. This is faster than weight quantization because on most hardware, INT8 computations are faster than FP32. But one challenge is we don't know the inputs of the neural network when we quantize the model. So, in order to determine the scale factors for each layer, we will need a calibration set that represents what kind of data we expect to see during inference time. If you ever come across the terms static or dynamic quantization, these refer to different ways of determining the scale factors of activations. If calibration is not done properly, you will encounter clipping in the network because the quantization is only able to handle floating points in a certain range, and anything outside of the range will clip to the max or min values.
To determine which type of quantization to use, it helps to look at the specifications of the hardware that you intend to do the inference on. Here I have pulled up the data sheet for the Nvidia A10 GPU, which is a popular choice for inference. According to the specification sheet, the FP32 performance of this GPU is 31 teraflops, whereas the INT8 performance is a lot faster at 250 tensor operations per second. This is thanks to its tensor core capabilities, but not all GPUs have this capability. So, on some older GPUs, you might find that the FP32 has the same performance as INT8.
One more thing that you should be aware of is the effect of outliers on quantization. One recent paper called LM INT8 found that in large language models with over 6 billion parameters, quantization doesn't work because of outlier features that cause the performance of the model to fall to close to zero. To understand why this is the case, consider what will happen if you have an outlier in the weight. What happens is the buckets become very large because there is only 256 buckets to cover all of the values between the minimum and maximum values, including the outlier. To solve this problem, they proposed a mixed decomposition scheme where the outliers are handled separately from the majority of the data. This is not necessary when you're running smaller models, but useful to know if you ever plan to quantize larger language models.
Now let's move on to the second method, pruning. The basic idea of pruning is you want to remove some of the connections in your neural network. This leaves you with what is called a sparse network, and in terms of the matrix computation, a lot of the values in the matrix get set to zero, which makes it cheaper to store and faster to compute. Once again, there are many different algorithms you can use to do your pruning, and in this video, I will only talk about the simplest one, magnitude pruning. In magnitude pruning, you first pick a pruning factor X, which denotes what proportion of the connections you would like to remove. Then, in each layer of the network, you set the lowest X% of the weights by absolute value to zero. The idea being that the lowest weights by absolute value, so the ones closest to zero, are the least important for the network to function. By removing some of the connections, your model will experience some degradation in accuracy. So, as an optional third step, you may want to retrain your model for a few more iterations while keeping the removed weights fixed at zero, and this is to recover some of the accuracy.
Now, it's important to note that just setting some of the matrix values to zero doesn't actually save space or make it go any faster because zeros take just as much space to store and just as much time to process as non-zero values. So, if you're doing pruning, you need to combine that with some sort of sparse execution engine that can take advantage of a specified neural network structure. Let me give you an example of what I mean by this. In general, when your GPU performs matrix multiplication, it iterates over slices of your two matrices, and for each pair of slices, it accumulates an outer product matrix, and the sum of all of this is the matrix multiplication. But even if you have zeros in the slices, it does not affect how long this operation will take. Compare this, on the other hand, with an algorithm that's specifically designed to multiply sparse matrices. The sparse matrix multiplication algorithm has a special trick that skips over all of the zero entries in a vector, so that the more zeros you have in the matrix, the faster the multiplication will be.
The last thing I will talk about is structured pruning. If you simply remove connections from a network without any further pattern, that is called unstructured pruning. But structured pruning is when you enforce more structure on which weights you are allowed to set to zero. One type of structured pruning is a two-out-of-four structured sparsity pattern. What this means is for each block of four consecutive matrix values, only two of them are allowed to be non-zero, and this allows you to store the matrix in a compressed format where only the non-zero values are stored along with indices for which values are represented in which positions as well. Nvidia's Tensor Core GPUs are able to execute this type of structured sparsity with greater efficiency. So, we see that for pruning neural networks, we need to design the pruning algorithm with the hardware in mind. Which pruning algorithm you should use will depend on which type of sparsity runs fast on the hardware that you intend to deploy your neural network on.
The third method of making our model more efficient is knowledge distillation, or sometimes called model distillation. So, what is knowledge distillation? In knowledge distillation, we first use the data to train a teacher network. After the teacher network has been trained, we then start training the student network to predict the outputs of the teacher network. Well, you might ask, why is it more helpful to have the student network predict the outputs of the teacher network instead of just training the student network from the labels? And the reason is basically the output of the teacher network contains more information, so it is faster and easier for the student network to learn from it. Assuming you're doing some kind of classification model, then the training data only has one label per training instance, but the output of the teacher network gives you a probability distribution over all possible labels, which is a lot more information to learn from.
Knowledge distillation has several advantages and disadvantages compared to other methods of optimizing your model. One advantage of knowledge distillation is you can modify the architecture of the student model to be different from the teacher model. For example, if your teacher model has 12 Transformer layers, that doesn't mean your student model has to have 12 Transformer layers and might have six or two or something like that. And this sort of architectural change is not really possible with quantization or pruning. Therefore, knowledge distillation has the biggest potential gain in speed compared to all of the other methods that we've seen. But the disadvantage is it's relatively more difficult to set it up because you need to set up the training data, which might be billions of tokens, and if the teacher model is a big model, then running inference over it can be a challenge. So, overall, knowledge distillation is relatively expensive. In my previous experience, this takes maybe 5 to 10% of the total compute or GPU hours needed to train the teacher model from scratch.
Here is one example. This too is a model trained with knowledge distillation, where BERT is a teacher model. In this model, they reduced the size of the BERT base model by 40% while retaining 97% of its accuracy, and the authors tell us how many GPUs and for how long they had to train this model. The student was trained on 8 GPUs for about 90 hours, so in total about 700 hours of GPU time. In comparison, the RoBERTa model, which is similar to the BERT model, required one day of training on a thousand GPUs, which is about 24,000 hours of GPU time, or around 20 times bigger. So, we see that in the distilled BERT example, training a model using knowledge distillation is a lot faster than training from scratch, but still requires a significant amount of compute.
The last category of optimizations are what I will group together and call them all engineering optimizations. At some point, you need to decide whether you want to run your model on CPUs or GPUs. In either case, making it run efficiently requires doing some integration between the hardware and the software. What I mean by that is your hardware might have the physical capability of running a model quickly, but at the same time, the software needs to know how to use the hardware capabilities, for example, vectorized operations to multiply large matrices in a parallel manner. GPUs are of course very good at this, but CPUs actually can do vectorized operations as well using some of the newer instruction sets like AVX2 and AVX512. Newer CPUs and GPU models have the ability to perform reduced precision and mixed precision operations faster than full precision, like in INT8 format, and this is useful for running inference on quantized models quickly as well. Some GPUs have the hardware capability to run sparse kernels, which is necessary to have a gain in speed when running pruned neural networks.
Another type of optimization is fused kernels. For example, PyTorch has a function called scaled dot product attention, and what this function does is it combines all of these operations that are typically seen together in a Transformer architecture, but it does it very quickly. In a Transformer architecture, we often do a sequence of operations in a fixed order, for example, multiply the query and key matrices, and then take a softmax, take a square root, and then apply dropout. And if we combine all of these operations into one single operation that is executed on the GPU, then this is a lot faster than if we executed each instruction in sequence. One popular way of implementing this is called FlashAttention. Not only does it fuse together these operations, but FlashAttention also does some tiling and some optimization according to the GPU's memory hierarchy to further reduce the amount of time needed to perform this operation. And you can see on the chart on the right here that the fused FlashAttention is a lot faster than naively doing all of the operations sequentially in PyTorch.
All of this might sound a little bit overwhelming, but really, it's not that complicated in practice because all you have to do in practice is convert your model that you have trained into some format that is executable by an inference engine that is optimized for whatever hardware that you intend to deploy on. The reason why you often need to use separate frameworks for training and inference is because the requirements for training a neural network is often quite different from the requirements during inference. When you're training a model, you need a library that can do things that are relevant during training, like loading the data from disk, processing it, doing gradient descent and back propagation, running evaluation, saving checkpoints, and so on. But none of that is really required during inference. When you're deploying a model for inference, however, the requirements tend to be quite different. The model needs to be small and fast, and needs to run efficiently on hardware that's probably different from what you trained the model on. So, it is often better to use a different library for inference. Two of the most popular libraries for inference are ONNX Runtime, which can run models that are stored in the ONNX format on a variety of different hardware, and another one is TensorFlow Lite if you prefer the TensorFlow ecosystem.
Let's summarize what we have covered in this video so far. First, quantization. Quantization uses less precise data formats to reduce the model size and latency. When you're reducing the format from FP32 to INT8, this results in a reduction of 4X. It is best used in combination with a reduced precision execution engine that is able to execute reduced precision formats faster, and a drawback is it can potentially result in a loss of accuracy, although hopefully not too much. Pruning is setting some of the weights of neural networks to zero to save space and compute, and in order for this to work at all, it requires an execution engine that's capable of executing sparse neural networks. And similar to quantization, it can potentially result in a loss of accuracy. Knowledge distillation is the only method we covered where you're able to modify the model's architecture, so the impact of this is varied depending on how you modify the architecture, but can potentially be much larger than any other method. The downside of knowledge distillation is it's relatively expensive to train. And finally, engineering optimizations. These should be used in combination with all of the above methods, and you should expect no loss in accuracy when employing engineering optimizations because the output should be identical.
Ultimately, all of these methods make a trade-off between development cost, inference cost, and model accuracy. Quantization and, to some extent, model pruning are two ways of reducing the model's latency and inference cost without being too difficult. But for both of them, you will potentially incur a slight loss in model accuracy. Knowledge distillation has potential to reduce your model size a lot further, but it is also more complicated and expensive to train, especially for larger models that are trained with lots of GPUs.
Thank you for watching, and I hope you will use these techniques in your own projects. If you found this video helpful, please don't forget to like and subscribe to my channel and get notified when I make new and helpful machine learning related videos. It will help me out a lot. Goodbye.