diff --git a/python/mlx/nn/layers/embedding.py b/python/mlx/nn/layers/embedding.py index 891627be..e05cfb5f 100644 --- a/python/mlx/nn/layers/embedding.py +++ b/python/mlx/nn/layers/embedding.py @@ -45,6 +45,9 @@ class Embedding(Module): group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = "affine", + quantize_input: bool = False, ): """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer.""" + if quantize_input: + raise ValueError("Quantized input is not supported.") return QuantizedEmbedding.from_embedding(self, group_size, bits, mode) diff --git a/python/mlx/nn/layers/linear.py b/python/mlx/nn/layers/linear.py index 8f00ff18..7b868b93 100644 --- a/python/mlx/nn/layers/linear.py +++ b/python/mlx/nn/layers/linear.py @@ -5,7 +5,7 @@ from typing import Any, Optional import mlx.core as mx from mlx.nn.layers.base import Module -from mlx.nn.layers.quantized import QuantizedLinear +from mlx.nn.layers.quantized import QQLinear, QuantizedLinear class Identity(Module): @@ -75,8 +75,36 @@ class Linear(Module): group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = "affine", + quantize_input: bool = False, ): - """Return a :obj:`QuantizedLinear` layer that approximates this layer.""" + """Return a quantized approximation of this layer. + + If ``quantize_input`` is ``False``, returns a :obj:`QuantizedLinear` + (weights are quantized). If ``quantize_input`` is ``True``, returns + a :obj:`QQLinear` (weights and activations are quantized). + + Args: + group_size (Optional[int]): The quantization group size (see + :func:`mlx.core.quantize`). Default: ``None``. + bits (Optional[int]): The number of bits per parameter (see + :func:`mlx.core.quantize`). Default: ``None``. + mode (str): The quantization method to use (see + :func:`mlx.core.quantize`). Default: ``"affine"``. + quantize_input (bool): Whether to quantize input. Default: ``False``. + + Returns: + QuantizedLinear or QQLinear: A quantized version of this layer. + + Notes: + Quantized input is only supported for ``"nvfp4"`` and ``"mxfp8"`` + modes. + """ + if quantize_input: + if mode not in ["nvfp4", "mxfp8"]: + raise ValueError( + f"Quantized activations are only supported for 'nvfp4' and 'mxfp8' modes, got {mode}." + ) + return QQLinear.from_linear(self, group_size, bits, mode) return QuantizedLinear.from_linear(self, group_size, bits, mode) diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py index 15bbfa76..57e7c888 100644 --- a/python/mlx/nn/layers/quantized.py +++ b/python/mlx/nn/layers/quantized.py @@ -25,13 +25,18 @@ def quantize( bits: int = None, *, mode: str = "affine", + quantize_input: bool = False, class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = None, ): """Quantize the sub-modules of a module according to a predicate. - By default all layers that define a ``to_quantized(group_size, bits)`` - method will be quantized. Both :obj:`Linear` and :obj:`Embedding` layers - will be quantized. Note also, the module is updated in-place. + By default all layers that define a ``to_quantized()`` method will be + quantized. Both :obj:`Linear` and :obj:`Embedding` layers will be + quantized. The module is updated in-place. + + Note: + ``quantize_input=True`` is only supported for ``"nvfp4"`` and ``"mxfp8"`` + modes and :obj:`Linear` layers. Args: model (mlx.nn.Module): The model whose leaf modules may be quantized. @@ -41,12 +46,23 @@ def quantize( :func:`mlx.core.quantize`). Default: ``None``. mode (str): The quantization method to use (see :func:`mlx.core.quantize`). Default: ``"affine"``. + quantize_input (bool): Whether to quantize activations. Default: ``False``. class_predicate (Optional[Callable]): A callable which receives the - :obj:`Module` path and :obj:`Module` itself and returns ``True`` or a - dict of params for `to_quantized` if it should be quantized and - ``False`` otherwise. If ``None``, then all layers that define a - ``to_quantized(group_size, bits)`` method are quantized. - Default: ``None``. + :obj:`Module` path and :obj:`Module` itself and returns ``True`` or a + dict of params for ``to_quantized`` if it should be quantized and + ``False`` otherwise. If ``None``, then all layers that define a + ``to_quantized()`` method are quantized. Default: ``None``. + + Example: + Weight only quantization for all layers that define a ``to_quantized()`` method: + + >>> import mlx.nn as nn + >>> nn.quantize(model, group_size=64, bits=4, mode="affine") + + Weight and input quantization for all linear layers: + + >>> predicate = lambda p, m: isinstance(m, nn.Linear) + >>> nn.quantize(model, mode="nvfp4", quantize_input=True, class_predicate=predicate) """ class_predicate = class_predicate or (lambda _, m: hasattr(m, "to_quantized")) @@ -54,8 +70,15 @@ def quantize( if bool_or_params := class_predicate(path, m): if hasattr(m, "to_quantized"): if isinstance(bool_or_params, bool): - return m.to_quantized(group_size=group_size, bits=bits, mode=mode) + kwargs = {"group_size": group_size, "bits": bits, "mode": mode} + if quantize_input: + kwargs["quantize_input"] = quantize_input + return m.to_quantized(**kwargs) elif isinstance(bool_or_params, dict): + if ("quantize_input" in bool_or_params) and not bool_or_params[ + "quantize_input" + ]: + bool_or_params.pop("quantize_input") return m.to_quantized(**bool_or_params) else: raise ValueError( diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index b259da3e..ebdfe580 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -204,6 +204,27 @@ class TestBase(mlx_tests.MLXTestCase): self.assertTrue(isinstance(m.layers[2], nn.QuantizedLinear)) self.assertTrue(isinstance(m.layers[2].scales, mx.array)) + m = nn.Sequential( + nn.Embedding(5, 256), nn.ReLU(), nn.Linear(256, 256, bias=False) + ) + nn.quantize( + m, + group_size=32, + mode="mxfp8", + quantize_input=True, + class_predicate=lambda path, module: isinstance(module, nn.Linear), + ) + self.assertTrue(isinstance(m.layers[0], nn.Embedding)) + self.assertTrue(isinstance(m.layers[1], nn.ReLU)) + self.assertTrue(isinstance(m.layers[2], nn.QQLinear)) + + # Check that Embedding does not support quantize_input + m = nn.Sequential( + nn.Embedding(5, 256), nn.ReLU(), nn.Linear(256, 256, bias=False) + ) + with self.assertRaises(ValueError) as context: + nn.quantize(m, group_size=32, mode="mxfp8", quantize_input=True) + def test_quantize_freeze(self): lin = nn.Linear(512, 512) qlin = lin.to_quantized()