Quantize module to QQLinear

[nastya236]

Extended Linear.to_quantize(..) and nn.quantize() to support quantization to QQLinear.

Added a flag quantize_input that defines the behaviour (default is False, so QuantizedLinear is used).

class DuckModel(nn.Module):
    def __init__(self, dim, n_layers):
        super(DuckModel, self).__init__()
        self.block = nn.Sequential(
            nn.Linear(dim, dim, bias=False),
            nn.ReLU(),
            nn.Linear(dim, dim, bias=False),
            nn.ReLU()
        )

        self.layers = [
            self.block for _ in range(n_layers)
        ]

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

model = DuckModel(128, 4)
x = mx.random.uniform(shape=(128, 128))
y_hat = model(x) 
nn.quantize(model, quantize_input=True, mode='mxfp8', class_predicate=lambda path, module: isinstance(module, nn.Linear))
y = model(x) 
mx.allclose(y, y_hat, atol=1e-3)

[angeloskath]

This is very nice.

What I am thinking is whether we should have quantize_input as a necessary argument or not.

So far all arguments are mandatory. When mode was added that means that every module that needs to work with nn.quantize needs to accept a mode argument.

We could do the same for quantize_input or we could make it optional and nn.quantize only passes it if it is True.

[awni]

We could do the same for quantize_input or we could make it optional and nn.quantize only passes it if it is True.

For the quantize-able modules we have now Embedding, Linear, SwitchLinear, it makes sense for all of them to support quantized_input. (For the embedding it would only be used when you call it as a linear layer Embedding.as_linear). So I'm ok with the option of just passing to all modules (but maybe only if it's True to avoid breaking the API for others).

If we don't intend to have all modules accept the parameter then ideally we should generalize the is_embedding check. Some options there are:

• Inspect the signature for the kwarg quantize_input. Hacky.. but maybe not too bad.
• Have a different method to_qq (or something) and check for it's existence

[angeloskath]

Hadn't thought of Embedding.as_linear at all. Good point.

I think passing it all the time makes sense. We did break the api for mode it doesn't seem that this will keep happening so let's do it.

[nastya236]

For the embedding it would only be used when you call it as a linear layer Embedding.as_linear

Good point. Do you think that it is needed? I had a feeling that typically embedding stays in bf16 for convergence. If we do want to quantize it, then we probably need either:

• A new module (something like QQEmbedding) that uses qqmm for as_linear. And we will change Embedding.to_quantize(quantize_input= quantize_input)
• Alternatively, we can use quantize_input in Embedding.as_linear directly, but it feels wrong..
• We skip embedding layers if quantize_input=True (at least for now).

Did you have some other option in mind?

[awni]

We skip embedding layers in quantize_input=True (at least for now).

I think that option is perfectly fine for now.

The main question is if we intend to require modules handle the argument quantize_input themselves in which case it's on the user to make sure they only request certain modules to use quantize_input, or do we choose for the user by only passing it to certain modules (e.g. nn.Linear).

I maybe just slightly prefer the first version. So in that version nn.Embedding.to_quantized should take quantize_input. And maybe at some point we implement QQEmbedding but also fine to just raise an exception if it's True for now.

The alternative would be to only pass quantize_input to modules that we choose. It's more opinionated and less flexible. But also probably easier to use for now.

I'm really okay with either way... we can always adapt later. It's quite a new feature so it's ok to be flexible based on how it's used.

[awni]

This is very nice! Minor request: would you mind adding a simple test in this file to make sure that quantizing a model with quantize_input produces QQ layers as expected?

[awni]

Looks good to go other than missing test! I'll run the tests after you add a test and then we can merge when they clear.