Skip non-quantizable modules in nn.quantize with a custom predicate

[shubhxho]

Fixes #3865.

With the default predicate, nn.quantize skips modules that don't define to_quantized(). Passing a custom class_predicate loses that behavior: if it returns truthy for such a module, the call raises Unable to quantize model of type ... and the whole model load fails.

class Gate(nn.Module):        # has params, no to_quantized()
    def __init__(self):
        super().__init__()
        self.weight = mx.zeros((8, 64))

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(64, 64)
        self.gate = Gate()

nn.quantize(Block(), group_size=64, bits=4)                                  # ok
nn.quantize(Block(), group_size=64, bits=4, class_predicate=lambda p, m: True)  # raises

This is why call sites like mlx_lm start every predicate with if not hasattr(m, "to_quantized"): return False.

The fix checks for to_quantized() before calling the predicate, so a module without it is skipped whether the predicate is custom or the default. Standard layers all define to_quantized(), so existing calls are unaffected. Updated the docstring and added test coverage for a custom predicate returning True and a params dict over a non-quantizable module.

[Pablosinyores]

Reproduced the reported failure on current main. Because quantize walks leaf_modules(), any leaf that does not define to_quantized reaches the predicate, so a broad predicate hits the raise on ordinary layers:

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(64, 64)
        self.norm = nn.LayerNorm(64)

nn.quantize(Block(), group_size=64, bits=4, class_predicate=lambda p, m: True)
# ValueError: Unable to quantize model of type <class 'mlx.nn.layers.normalization.LayerNorm'>

leaf_modules() returns ['proj', 'norm'], so this is not an exotic case, any model with a norm or activation layer trips it. Skipping those rather than raising looks like the right direction to me, and doing the hasattr check before calling the predicate reads cleaner than the nested version it replaces. black is clean on both files.

One thing I would want a maintainer to weigh explicitly: the change removes the Unable to quantize model of type ... error entirely, so a predicate that deliberately selects a type that cannot be quantized now silently does nothing. class_predicate=lambda p, m: isinstance(m, nn.LayerNorm) used to say why nothing happened; after this it returns a model that is simply unquantized, with no signal.

Skipping incidental leaves and staying silent when the user asked for something impossible are two different things, and this collapses them into one behavior. If you want to keep a signal without bringing back the old over-eager raise, one option is to raise only when the predicate returned true for at least one module and nothing was quantized at all. Happy either way, but it seems worth deciding rather than dropping the diagnostic as a side effect.

Also worth updating the docstring line about the default: with class_predicate now defaulting to lambda _, m: True, the "all layers that define a to_quantized() method are quantized" sentence is still accurate, but it now describes the hasattr guard rather than the default predicate. The added sentence covers it, just noting the two are doing different jobs now.

[nastya236]

Thank you for flagging this! Will look at it today

[nastya236]

Thanks for looking into this and for the detailed issue.

That said, I think the current behaviour is intended:

• If class_predicate is not provided -- quantize every module that defines to_quantized(), and skip the rest.
• If class_predicate is provided -- quantize every module the predicate selects (True or a params dict). If a selected module does not define to_quantized(), we raise. This is deliberate: it guarantees that anything explicitly selected for quantization actually gets quantized, rather than being silently skipped.

In other words, class_predicate is an explicit selection, not a filter over "quantizable" modules. Silently skipping a selected by a predicate but unquantizable module would hide real mistakes, for example a custom layer meant to be quantizable but missing to_quantized().

If you want to quantize a custom module, the intended way is to define a to_quantized() method. And if you specifically want a broad predicate to skip non-quantizable modules, you can compose that yourself:

nn.quantize(model, class_predicate=lambda p, m: hasattr(m, "to_quantized") and my_predicate(p, m))

Given that, I don't think we should change this behaviour, so I'll close the PR.