Brevitas Quantization Library - Pablo Monteagudo Lago, AMD

This talk covers Brevitas, a neural network quantization framework in PyTorch developed by AMD. The speaker discusses the necessity of quantization for enhancing model efficiency, explaining how it reduces the size and computational requirements while maintaining performance. Brevitas facilitates both quantization-aware training and post-training quantization, and it is designed with flexibility and composability in mind for researchers. The talk highlights specific techniques, such as the Croq algorithm, error correction methods, and an illustrative pipeline for quantizing a model from Hugging Face. Brevitas integrates with other frameworks, supporting various deployment targets and enabling users to fine-tune quantization parameters easily to optimize performance.

Hello, I'm Pablo Monteagudo Lago, and I'm going to talk about Brevitas, a neural network quantization framework in PyTorch from AMD. And the central question of this talk is, how does quantization tooling need to look like to support both cutting-edge research and multi-target deployment? I would argue that the answer is to be both flexible and composable, and I will show you how Brevitas delivers both of these features.

So, a little bit about the road map of this of this presentation. I'm going to first talk about about the motivation for quantization, try to answer the question, why do we even care about quantization? Go on to briefly introduce Brevitas and how it fits into the quantization landscape. Then we'll talk about flexibility and composability and why these are features that we require that we want in a

quantization tooling to support cutting-edge research, and talk about an specific instance of of research done with Brevitas called Croq. Finally, to finalize, we'll go through an end-to-end example quantization pipeline in Brevitas in which we'll start from a pre-trained model, apply several algorithms in it to mitigate the quantization error, and finally deploy to our favorite deployment target. So, why quantization is important? The main key observation in deep

learning is that model quality quality usually scales with size. But the model size is not only the the only driver of cost because for language models, when we spend more more compute time, we also increase we can through techniques we when increase the the inference the inference time through techniques such as chain of thought, we can also improve the reasoning capabilities. So, the So, the inference cost

is growing from two directions. We have bigger models, and we have longer generations. So, the So, the problem is not yet more and now inference cost is a first-class concern and we want to trade off and we want to improve the trade off between inference cost and model quality by addressing these two concerns. So, this is where quantization comes comes into play. Quantization is an useful tool

to And what is quantization? So, quantization is the process of mapping the weights and activations of a model from high precision either like BFloat16, FP32 to a lower precision representation. For example, INT8 or MXFP4. So, we'll see like illustrative example on how uh a simple quantization pipeline would look like. So, first we start with the original distribution of our weights that is like with weights all over

the the real line and can take any value. And now we need to quantize it to a discrete grid. So, first what we what we do is a scale the weights a scale the so it matches the the grid. In this case, indicated by the by the green dashed lines. After this, what we can do is round the weights to the nearest point in the in this

grid and this is indicated by the arrows in which we're mapping the And finally we des scale back, we apply the scale factor to map to map to the original range. And but this process of mapping to the closest point at this degree, we can see that can yield a catastrophic accuracy degradation when we are in in with data type that have lowest resistivity. And we'll see

how algorithms come into play error. So, what are the benefits of doing this? What are the benefits of mapping to these lower precision representations? And the benefits depend on the hardware and software platform, but can generally fall into four categories. Memory, when we go from 16-bit to 4-bits, we are using four times less storage and we are the arithmetic in low precision, our arithmetic units are smaller

and we see and we see improvements in terms of power, use less power, the throughput increases and the area of these low precision arithmetic units in the chip is smaller and we can have more of them or use the this extra area for things. So, it's clear that quantization is important and how do we enable in practice? And precisely Brevitas is a library to do so that

has been around for for a while, eight years, with over five 5,000 5x downloads and 1.5k GitHub stars. And one Brevitas is that it targets the full AMD product range going from CPU, GPU, MPU and FPGA. And this library mimics the standard interface of PyTorch and the modules and supports both quantization aware training in which our starting point is a random model that we want to train

size at the same time and post training quantization model in which we start from a model that has already been pre-trained pre-trained and we want to map map it to our deployment target. And Brevitas one of the key principles of the construct of the software architecture in Brevitas is that it is built for composability and flexibility our quantization research. And to cite a few works, I'm going

some of them are mixed one, Cronus and accumulator aware quantization. And I'll be talking a little bit about Cronus and how it how it improves the state of the art in error correction algorithms. And finally, Brevitas is not is not meant to be used in isolation. In integration with hugging face transformers, you can load your favorite model and then after you are happy with your with your

results after quantizing the model and apply algorithms on top of it, you might want to export it to your favorite framework such as VLLM, GUF, llama.cpp, etc. So, the main thing that I want to highlight is for both researchers and practitioners. And it spans the full pipeline going from experimentation that we when we are trying to apply algorithms to maximize accuracy to deployment when we have a

target in mind in which we want to deploy our quantized model. And how Gravitas fits into the So, quantization involves many design choices. We have the data types we can quantize to integers, to floating point, mini floats. We have the quantization parameters such as a scale and a zero point, and we need to tune them to maximize the the how how well we fit our our to

our The algorithms that we apply on top to maximize accuracy. And all of these can be tweaked to maximize performance. So, this is space of combination grows combinatorially. But then when we take a look at most of most of the other quantization tools, most of them focus on a few fixed recipes. It makes sense when you are focusing on production performance for a specific target. But when

you are a researcher, you want to combine the different pieces in unusual ways to get your best the best results. And this is what we focus in Gravitas. We want to navigate this space and support any sensible combination of these different design choices. And for that, we need a framework that is both flexible and composable. I will talk a little bit how we rate Gravitas is built

around these principles. So, I've I've mentioned before that there are many knobs in quantization, and to emphasize on that point, this is a a fake quantization formula type in which we are projecting our weight to a to a discrete integer grid. And each each color-coded term represents a design decision. So in red, we have the we have the design choices regarding to the precision, integer, binary, our

data type binary, floating point. If we are using floating points, we might choose to represent the the infinite, the NaNs, or we may even have like custom formats. For example, we might use different bits for for the a different bits for the exponent and mantissa compared to a standard data types. So there are Ideally, we would like to starting with a with an existing quantizer, we might

be we might want to tweak any of these, which is what I'm showing in this code listing. In which basically, we start with a with a quantizer, and then we we change the parameters that we are interested in. And this can be done without needing to subclass anything fancy. Then we have some choices regarding the scale and zero point because these quantization parameters that that allow us

to move the grids and and scale them can have different granularities. They can be shared shared for all the weights in the tensor, for all the weights in a channel, we subgroups within a channel. And again, like this can all easily be tweaked. finally, we want to apply the algorithms because this might mapping mapping directly to the grid might not be enough and might cause like quality

degradation. I mean, we might want to to fix this. And in Brevitas, we have several entry points. We have entry points for quantizing the fusion models. classification classification image classification models. And we have an LLM entry point, which is the the CLI example that I'm showing here, in which you can select the hyper parameters of Um as a working example for some somebody who does research in

in quantization, I'm going to talk about a little bit about gradients or against. if we think of a rounding function, the rounding function is a stepwise function and its gradient is which is almost everywhere. So, if if we were back propagating through this function, the gradient would be vanishing everywhere and we wouldn't be getting any any updates on our weights. So, we need to define a meaningful

gradient for it and the usual choice is to assign the gradient of this rounding function to one. Thus, as for our back work this function is behaving as an identity. But the problem with this is that it misses the periodic structure of rounding. Um some researchers from Microsoft identified this limitation and proposed a Fourier-based surrogate that captured this structure. But if you are a researcher, you want

to focus on your innovation. You don't want to rewrite any boiler to hook this innovation to a minifloat quantizer and integer quantizer. You want to just focus on your idea, iterate on it and be able to run any experiments easily. So, because this innovation only affects the rounding model and the rest of the quantization pipeline remains unchanged, in Brevitas what we do is allow the the researcher

to implement his own the new rounding model and then this new rounding model, which incorporates the custom the custom gradient function, can be hooked into any existing quantizer. And by any, I mean literally any any quantizer, which makes experimentation easier because function has has some several hyper-parameters and tweaking them can be done in a very transparent way. So, as a researcher, Brevitas allows you to only focus

on the new rounding logic while the integration is handled handled transparently for the for the researcher. So, I have to talk about flexibility to talk about composability in the context of PTQ algorithms. As I was saying befo- before, mapping mapping weights to a low precision grid by directly rounding to the ne- to the nearest one can be suboptimal and can yield like catastrophic degradation. And post-training quantization

address this in two complementary stages. And the following the the following classification that I'm going to do is a is one that we have, but there are many, but we find that it's it is a useful one to think about composition of PTQ techniques. So, the first uh stage that about is the transform stage. And this stage happens before we introduce any quantization noise in our model.

This is a stage that aims to answer the question, how can I transform my model in a way that mitigates the quantization So, for example, outliers are especially harmful for when quantizing for quantization. Uh they take a lot of their dynamic range. So, for example, small weights might get clamped to zero because of all because these outliers are moving are stretching out the the quantization So, ideally,

we don't want outliers in our weights and activation distributions. And a technique to address this is to rotate this weight vector in a way that after rotations, we end up with distributions that are more Gaussian shape and do not have these heavy tails that are harmful for quantization. And in in here, we are altering these weights and activation distribution to make them more amenable for quantization. And

the next stage corresponds that answer the question, how can I map these transformed weights uh to my discrete grid and recover from quantization error because we we don't to to assign the weights to the nearest point in the grid. We want to assign our weights selectively in a way that we are minimizing for output error at the at the output of a one layer. And there are

many algorithms for doing this. Kronos, AutoRound. I'll be expanding on Kronos and Kronos after So, when we think of a modern PTQ pipeline, we need to think about these two stages and that we are first going to transform our model and then commit to the commit to the round specific rounding choices. And once we are happy with the with the performance, we can export our model to

a serving framework such as a VLLM. So, let's put some numbers to emphasize what why algorithm composition is important. And imagine you are searching and you got a novel data type. In this case, it's a variation of MXFP4, but with two bits and integer data underlying integer data type. So, if we were quantizing quantizing uh this Llama Llama 1B model to to this a new data type

by just round to nearest round to the nearest point in this one. We got a perplexity, which is the accuracy metric that we are interested in, that is on the order of millions. So, basically our model has has turned completely unusable. So, we might think, okay, the my problem is that my data type is not expressive enough and what I going to do is go I'm going

to relax the constraint on my data type. I'm size and use a more expressive scale. But, I'm going to just round to nearest. And what I see is that the again, the result is in the order of the millions. My model is completely unusable. So, if I instead change my mindset and instead of thinking of more expressive data types, I I start thinking about algorithms, how can

I apply algorithms I'm going to do it on my original data type, in the one which was supposedly not good enough, not expressive enough. And the first thing that I'm going to do is apply rotations. I'm going to at the input of one layer and on the weights I'm going to rotate them and this this technique is called Hadamard incoherent processing. So, I get a notable recovery

even with the coarser with the coarser scales. Still, the result is in the order of the thousand, so pretty far away from our from our floating point baseline, but still even in in in the coarser data type we were able to bring the we were able to bring our quality with we were not able before when we turned to a more expressive data type. And finally, if

we think of okay, I'm going to after I suppress the layers I'm going to apply a technique that allows me to selectively assign my weights to their floor or ceiling and minimize the and minimize the quantization error. And by doing so, I get a perplexity and which you might is it is still like far away from the original baseline, but I could only get to this perplexity

by refining my algorithmic tooling and not like moving on to moving to more expressive data types. And precisely, Gravitas allows you to combine these in a single reproducible pipeline and transparently like combine them and use them until you get to your desired accuracy metric. So, now I'll talk a little bit about a research done with Gravitas. And I'm a Kronos, which is an algorithm like the rounding

choices for a given weight. So, the thing So, the thing is if I think of a weight of vectors, what I'm going to do is I'm going to quantize the weights one by one in a greedy fashion. I start with one and then which is called the greedy quantization step in which I'm going to assign in first weight to a given This is the output at the

error at the the that minimize the error in the output. And after quantizing this weight, there's some quantization error that has been incurred. And what I'm going to do is adjust the remaining weights to compensate for this error. So, my hope is that when I quantize the next one, already compensated for the for the error from before. And And the idea is this is a a greedy

a greedy way to And the difference between GPTQ, which was uh which is a well-known algorithm for doing this, and Chronos, is that while GPTQ only focus minimizing the error at the output of the layer, what Chronos does is also take into account what happened before. We are also not only trying to compensate for the error that happened within this particular layer that we are uh quantizing,

but also what happened the error that came from the quantization of the previous layers, which is like to uh in this uh red-colored uh axis. Um Oh, yeah. And this extra this extra expressive in the objective only cuz when you are when you are trying to solve this least square problem efficiently, only causes like a small increase in quantization run time compared to GPTQ. So, how do

these two algorithms compare when we take a look at the results? And we see that Chronos outperforms other greedy error correction algorithms like GPTQ or GPFQ. An important thing is that this difference becomes more significantly when you move uh to lower bit widths, from 4 bits to 3 bits. And trying to because at each at each when you are quantizing one layer, you are not able to

correct all the error. When you when when you go to the next one, Chronos also tries to compensate for the error that comes from the previous ones. Um So, we can see that it outperforms all the the rest of the greedy quantization algorithms. We compose Kronos with other algorithms. What happens when we apply algorithms to remove remove outliers before commit applying Kronos to uh select our running

choices. And we are going to apply again in Hadamard Incoherent Processing. And what we see is that Kronos and Incoherent Hadamard Incoherent all the results in which we apply Kronos isolation and we get a result that is much closer to the original floating point uh floating point result even in three bits, which is quite impressive. So, with this what I want to highlight is that Brevitas algorithm

composition is key for performance and Brevitas Um, yeah. Um, finally I want to rush through a uh end-to-end like quantization pipeline in Brevitas to see a little bit how this maps to how this process that I was talking before maps to code. Um, this will be like the full code listing in which we start from a hugging face model. We apply apply rotation equalizations to mitigate the

outliers. We quantize the weights and activations. We apply Kronos for error correction and finally we deploy to VLLM. Um, and again, like we use rotations for for mitigating the outliers. An important thing is that we can apply these rotations in a way computational graph. So, we are what we are going to do is identify regions in the model in which if I rotate a set of weights,

I rotate by the transpose the other set of weights, the result if we were doing the computations in infinite precision would be the same. But but actually the resulting distributions are don't have these outliers that are especially And Brevitas is able to use Dynamo to trace the graph and identify these the regions that that verify this property. And this is in contrast with other libraries that require

modifying modifying the module manually and adding and making the user adding these rotations, which regard of where do these rotation equivalent regions live. And next, uh Gravitas offers different ways to um different ways to insert the quantization nodes ranging from manual in which you have full full control of the of the quantized layer you're using to programmatic and effects in which you placing modules by their quantized

counterparts. So, you can have full control or opt for something fully automatic depending on on your use case. And then we have the state two data correction in which we can apply Kronos or we can apply GPTQ and look this is complete is very similar. So, if you want to replicate a resource table like this, if you are a researcher, this can be this is really easy.

It's just choosing the right algorithm in a YAML config and just running these experiments. when you want to export them export the like the path from research to production, Gravitas supports multiple export targets via LLM, TIF, ONNX. And if you And therefore ensuring a smooth transition from experimentation to deployment. And if you are interested in another deployment target, please reach out. And finally, just a recap. And

what open source quantization library from AMD which offers a unified environment for modern PTQ. It's leverages the PyTorch latest PyTorch features, integrates with the PyTorch ecosystem, and and allows you to export to to multiple deployment targets. Gravitas, please let us know because we we welcome contributors and collaborators. Um thank you so much for your attention and if you have any question, please let me know. >> [applause]

>> Yep. Or can you Can you repeat the last week? Both both both of them are Yeah. If If you want to work with FX graph, you can. And if you want to work in eager mode, that you can. And yes, you replace the for example, the linear models by one linear and that can be done programmatically. So, either either either works. Yeah, definitely. So, so yeah.

Like if you have like you want It is It is that is not currently supported is very very easy to to add it. But for example, the export to to VLLM is Yeah, you can target any any backend within VLLM that doesn't need to be like AMD. Or if you want like you can export to ONNX and that can Yeah, you can export to many intermediate to

transit You transition from that intermediate representation to your backend doesn't need to be like AMD. Real It doesn't need to be AMD, let's say. Yeah. Yeah. Yeah. >> Yeah. Uh so so so in brain in theory like it should be because what what you what you can do is like you like the effects representation to like um to and basically take that as an unit of computation

and replace it by your counterpart quantized unit of computer. Uh it doesn't you don't have that out of the box and you would need to tweak a tweak a little bit, but it's definitely possible. If you have if you have your definition for the quantized your full quantized model. So so Brevitas is stack that it's it's not like it is like more far away from kernels and

so if you want to like map to a specific like backends or like kernels, you need to like have a an export target an export target for that. But Brevitas operates like when you are trying to I was saying of uh like maximizing accuracy, that happens usually like in in eager mode before you uh start making commitments to the specific targets that you want to the to

your specific Oh. Yeah. Thank you. Yeah, that definitely let's let's think afterwards in case somebody has more questions. Yeah, I I guess like the the idea is to be like very flexible for researchers because uh the the way they usual uh quantizers in PyTorch come, they they come as a monolithic block and for example if you you pick the scales factors like my no uh you might

not be happy with that and changing that requires like manually like uh my requires modifying the uh the module that Py Py PyTorch made or other like different choices in the in the quantization pipeline. And the idea is instead of thinking of quantizer as monolithic modules that you that you have and you have one quantizer for integers for data types, you you think of quantizers are as

I have this rounding function, I have this clip clipping function, I have this uh way to pick the scale and you you combine them as you want to create to instance of quantizers. So for our researchers as as I was saying, if you want to use a different backward function for for your model, this is as easy as I define this and when I when I create

my quantizer, I say my rounding function is this, instead of requiring you to subclass and make like any your custom model. And this might not And you would need to do that for an integer quantizer and floating point quantizer, which requires uh which requires writing my required writing many boilerplate code to actually like a streamline your experimentation. >> So, when when I quantize, for example, the model

to three or two bits and I want to serialize it, does it does ReBITAS support some back-end format so it is really compressed to to the bit width that I quantize the model? >> Uh no no, like [clears throat] the part the part of you packing the weights What ReBITAS does is that apply like fake quantization. So, you have you Most of the time you have your

original You think of quantization as you have the original weights and then you have a mapping that maps to the to the quantized weights. But then, how you get those quantized weights and pack them a 16-bit container with like four that represents four weights, that's not uh handled in ReBITAS. That would be like an export flow for a specific deployment Okay, thank you so much for your

attention.