[Feature] Support for out-of-source quantizers

[Giuseppe5]

Hello,

I am Giuseppe, one of the main maintainers of Brevitas.

I had the pleasure of chatting with @shimmyshimmer last week during the PyTorch conference, during which the topic of QAT and Torch AO integration came up.

I was curious to see if/how it could be extended to support other quantizers (such as Brevitas), and I believe it is fairly straightforward, although there are a few changes that are required to make this a bit easier.

Issues and possible solutions

The biggest issue is that the function that applies quantization at the moment is not easily overridable/modified, since it's a function call within a much bigger staticmethod of the FastLlama class.

I forked the repo, to propose a tentative solution to this problem. I am happy to accept other ideas and/or contribute with a PR, if that works with you.
These are the changes required:

Giuseppe5@58c22f0

Other smaller issues are related to the specialization around Torch AO naming scheme for quantizers (i.e, weight_fake_quantizer and activation_fake_quantizer).
There are easier out-of-source workarounds for this but maybe it can be abstracted to something more general?

Example

Starting from the original QAT notebook, I created a slightly modified one that works with Brevitas and my fork of unsloth.
You can find it here:

https://colab.research.google.com/drive/1HhetpDq3oKTN9VIeS3GCSWEWKi7PXG0r?usp=sharing

The main modifications are contained in a block called Brevitas quantization.
It is a very minimal examples, but it could be easily extended to other quantization formats.

What comes next

There are a few (minor) missing features compared to the current integration, like fusing back LoRA adapters into the weights. We believe this is easy to implement if everything else works as planned.

The main absence in the example above is the export pathways.

Brevitas can easily decouple quantization application from quantization representation, which means we can easily adapt and implement new export formats (for example, mimicking what Torch AO does, if that is what users want).

We currently provide several export formats (e.g., ONNX through optimum), and we are planning to expand to more (e.g., export to vLLM), but we would love to hear what in your opinion are the most useful export/serialization formats we should target.

cc @nickfraser

[Datta0]

Hey @Giuseppe5, if you have a working PR, I'd be happy to review and see what it looks like...
Also where can I read up more about brevitas and what is the core thing you're trying to solve/cater to as to get a better idea ?

[Giuseppe5]

I will open a PR soon and tag you, many thanks!

About Brevitas, I guess you can start from the (ever improving) documentation here: https://xilinx.github.io/brevitas/v0.12.1/
It's a research project that has been going on for almost 9 years now, so it's difficult to summarise it but I'll try to be as concise as I can.

Brevitas is a PyTorch quantizer extremely geared towards flexibility, in various forms:

• Possibility to use commonly used dtypes (Int8, FP8/FP4, MXFP), while also coming up and simulate the effect of your custom dtype, like done here.
• Possibility to finely customize the quantizers, for example deciding how to compute scale factors and zero point, as well as replicating and expanding on more complex scale/zp structures, like the GGUF's ones.
• It's possible to insert quantization by replacing layers with their quantized counterparts, or by obtaining an FX-graph and doing more fine-grained quantization surgery, or by defining your model from scratch with quant layers.
• In terms of algorithms, we support all the latest most important quantization algorithms (GPTQ, SpinQuant, Qronos, etc.), and it is possible to combine all the dtypes and quantizers mentioned above with all the quantization algorithms.
• We support PTQ and QAT and everything in between, and since we use the same classes to do both, there is no need to use special procedures or conversion to switch from one to the other.
• Finally, we decouple everything mentioned above from how we represent quantization, allowing us to support various export formats and allowing us (and any users) to add new ones if they want. Currently we support ONNX, torchscript (looking into if/how to support torch.export), GGUF and others.

In terms of users, we mainly target two types of users:

• Quantization researchers who want to have flexibility in implementing and verifying their new ideas
• Users who want to deploy their quantized models but they need to squeeze out as much accuracy as possible. For example, it is possible to do extremely heterogeneous quantization (using different bit-widths per layers, or using FP8/INT8 for some layers and MXFP8/MXFP4 for other layers).

Our goal for this integration would be to cater to both these groups.

In this context, we would like to hear from you what are the most important export paths after a model has been fine tuned through Unsloth, so that we can add support for those, once quantization has been added into the mix.

I hope this is equally concise and clear but please feel free to ask more questions if you have any, either here or by reaching out directly!