Tinygrad Quantization Support [WIP]

[KhanerX]

What I did:

1. Define custom layers for Affine Quantized models, including integer weights, float16 scales and biases (zero point correction)
2. Load MLX-Community quantized model and unpack the weights.
3. Write the forward logic for quantized layers, following this paper (see section 2.3)

Todo:

• [ ] write tests, test with multiple nodes and different llama models
• [x] support 4-bit quantization
• [ ] do the forward math in Integer (see section 2.2 of mentioned paper)

With this first commit, you can run exo --run-model="llama-3.2-1b-8bit" with tinygrad backend and "mlx-community" model.

[KhanerX]

Also, I'm doing math in float32 right now, which adds overhead. when I change it to float16, I think something overflows and model outputs nothing. I will fix this.

[AlexCheema]

This is a great start - I tested this and it works. That's awesome because that means we can support any mlx model in tinygrad.

Are you sending parameters in float32 to the GPU? Or are they being sent in fp8? Just wondering what kind of speed to expect here. Wonder how close this gets to MLX quantized

[varshith15]

tested it out on m3 pro for mlx-community/Llama-3.2-1B-Instruct-8bit

old pr:

new_pr:

mlx:

just 2.5x slower now 🚀🚀

@KhanerX @AlexCheema

[Evanev7]

Thanks so much for your contribution and for taking the time to open this PR.

Since this repository has been fully rewritten and the license has changed, I’m closing all existing open PRs to avoid confusion and to align with the new codebase.

I really appreciate your interest in the project, and you’re very welcome to open a new PR against the updated version if you’d like and we look forward to reviewing it!