Quantize module to QQLinear (#3106)

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)

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)
 
 

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(

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()