mlx 1.0.0

data/mlx/python/mlx/nn/layers/quantized.py

@@ -0,0 +1,426 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import math
+from typing import Callable, Optional, Union
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+from mlx.utils import tree_map_with_path
+
+
+def _defaults_for_mode(mode, group_size, bits):
+    mode_defaults = {
+        "affine": (64, 4),
+        "mxfp4": (32, 4),
+        "nvfp4": (16, 4),
+        "mxfp8": (32, 8),
+    }
+    default_group_size, default_bits = mode_defaults[mode]
+    return group_size or default_group_size, bits or default_bits
+
+
+def quantize(
+    model: Module,
+    group_size: int = None,
+    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()`` 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.
+        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 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()`` 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"))
+
+    def _maybe_quantize(path, m):
+        if bool_or_params := class_predicate(path, m):
+            if hasattr(m, "to_quantized"):
+                if isinstance(bool_or_params, bool):
+                    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(
+                        "``class_predicate`` must return a bool"
+                        " or a dict of parameters to pass to ``to_quantized``"
+                    )
+            else:
+                raise ValueError(f"Unable to quantize model of type {type(m)}")
+        else:
+            return m
+
+    leaves = model.leaf_modules()
+    leaves = tree_map_with_path(_maybe_quantize, leaves, is_leaf=Module.is_module)
+    model.update_modules(leaves)
+
+
+class QuantizedEmbedding(Module):
+    """The same as :obj:`Embedding` but with a  quantized weight matrix.
+
+    :obj:`QuantizedEmbedding` also provides a :meth:`from_embedding`
+    classmethod to convert embedding layers to :obj:`QuantizedEmbedding`
+    layers.
+
+    Args:
+        num_embeddings (int): How many possible discrete tokens can we embed.
+           Usually called the vocabulary size.
+        dims (int): The dimensionality of the embeddings.
+        group_size (Optional[int]): The group size to use for the quantized
+            weight. See :func:`~mlx.core.quantize`. Default: ``None``.
+        bits (Optional[int]): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
+        mode (str): The quantization method to use (see
+           :func:`mlx.core.quantize`). Default: ``"affine"``.
+    """
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        dims: int,
+        group_size: int = None,
+        bits: int = None,
+        mode: str = "affine",
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
+        self.mode = mode
+
+        # Initialize the quantized weight
+        scale = math.sqrt(1 / dims)
+        weight = mx.random.normal(shape=(num_embeddings, dims), scale=scale)
+        self.weight, self.scales, *biases = mx.quantize(
+            weight, group_size, bits, mode=mode
+        )
+        self.biases = biases[0] if biases else None
+        self.num_embeddings = num_embeddings
+        self.dims = dims
+
+        # Freeze this model's parameters
+        self.freeze()
+
+    def __call__(self, x):
+        biases = self.get("biases")
+        return mx.dequantize(
+            self["weight"][x],
+            scales=self["scales"][x],
+            biases=biases[x] if biases is not None else None,
+            group_size=self.group_size,
+            bits=self.bits,
+            mode=self.mode,
+        )
+
+    def as_linear(self, x):
+        """
+        Call the quantized embedding layer as a quantized linear layer.
+
+        Use this for example when input embedding and output projection
+        weights are tied.
+        """
+        return mx.quantized_matmul(
+            x,
+            self["weight"],
+            scales=self["scales"],
+            biases=self.get("biases"),
+            transpose=True,
+            group_size=self.group_size,
+            bits=self.bits,
+            mode=self.mode,
+        )
+
+    def _extra_repr(self):
+        return (
+            f"{self.num_embeddings}, {self.dims}, "
+            f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
+        )
+
+    @classmethod
+    def from_embedding(
+        cls,
+        embedding_layer: Module,
+        group_size: int = None,
+        bits: int = None,
+        mode: str = "affine",
+    ):
+        """Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer."""
+        embedding_dims, dims = embedding_layer.weight.shape
+        ql = cls(embedding_dims, dims, group_size, bits, mode=mode)
+        ql.weight, ql.scales, *biases = mx.quantize(
+            embedding_layer.weight,
+            group_size,
+            bits,
+            mode=mode,
+        )
+        ql.biases = biases[0] if biases else None
+        return ql
+
+
+class QuantizedLinear(Module):
+    """Applies an affine transformation to the input using a quantized weight matrix.
+
+    It is the quantized equivalent of :class:`mlx.nn.Linear`. For now its
+    parameters are frozen and will not be included in any gradient computation
+    but this will probably change in the future.
+
+    :obj:`QuantizedLinear` also provides a classmethod :meth:`from_linear` to
+    convert linear layers to :obj:`QuantizedLinear` layers.
+
+    Args:
+        input_dims (int): The dimensionality of the input features.
+        output_dims (int): The dimensionality of the output features.
+        bias (bool, optional): If set to ``False`` then the layer will not use
+            a bias. Default: ``True``.
+        group_size (Optional[int]): The group size to use for the quantized
+            weight. See :func:`~mlx.core.quantize`. Default: ``None``.
+        bits (Optional[int]): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
+        mode (str): The quantization method to use (see
+           :func:`mlx.core.quantize`). Default: ``"affine"``.
+    """
+
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        bias: bool = True,
+        group_size: int = None,
+        bits: int = None,
+        mode: str = "affine",
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
+        self.mode = mode
+
+        # Initialize the quantized weight
+        scale = math.sqrt(1 / input_dims)
+        weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(output_dims, input_dims),
+        )
+        self.weight, self.scales, *biases = mx.quantize(
+            weight, group_size, bits, mode=mode
+        )
+        self.biases = biases[0] if biases else None
+
+        # And bias if needed
+        if bias:
+            self.bias = mx.zeros((output_dims,))
+
+        # Freeze this model's parameters
+        self.freeze()
+
+    def _extra_repr(self):
+        out_dims, in_dims = self.weight.shape
+        in_dims = (in_dims * 32) // self.bits
+        return (
+            f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, "
+            f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
+        )
+
+    def __call__(self, x):
+        x = mx.quantized_matmul(
+            x,
+            self["weight"],
+            scales=self["scales"],
+            biases=self.get("biases"),
+            transpose=True,
+            group_size=self.group_size,
+            bits=self.bits,
+            mode=self.mode,
+        )
+        if "bias" in self:
+            x = x + self["bias"]
+        return x
+
+    @classmethod
+    def from_linear(
+        cls,
+        linear_layer: Module,
+        group_size: int = None,
+        bits: int = None,
+        mode: str = "affine",
+    ):
+        """Create a :obj:`QuantizedLinear` layer from a :obj:`Linear` layer."""
+        output_dims, input_dims = linear_layer.weight.shape
+        ql = cls(input_dims, output_dims, False, group_size, bits, mode=mode)
+        ql.weight, ql.scales, *biases = mx.quantize(
+            linear_layer.weight,
+            group_size,
+            bits,
+            mode=mode,
+        )
+        ql.biases = biases[0] if biases else None
+
+        if "bias" in linear_layer:
+            ql.bias = linear_layer.bias
+
+        return ql
+
+
+class QQLinear(Module):
+    """Quantizes the input and applies an affine transformation using quantized weights.
+
+    Two use cases are supported:
+
+    1) **Eval**:  The weights are frozen and stored in quantized form together with
+       their scales (``self.weight`` is quantized and ``self.scales`` is provided).
+    2) **Train**: The weights are stored in higher precision and are quantized on
+         the fly during computation so that gradients with respect to the weights
+         can be computed.
+
+    To switch between the two cases, use ``layer.eval()`` and ``layer.train()`` respectively.
+
+    Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer
+    quantizes the input as well and includes weights in gradient computations.
+
+    :obj:`QQLinear` also provides the class method :meth:`from_linear` to
+    convert :class:`mlx.nn.Linear` layers to :obj:`QQLinear` layers.
+
+    Note: This layer does not support a bias term yet.
+
+    Args:
+        input_dims (int): The dimensionality of the input features.
+        output_dims (int): The dimensionality of the output features.
+        group_size (Optional[int]): The group size to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
+        bits (Optional[int]): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
+        mode (Optional[str]): The quantization method to use (see
+            :func:`mlx.core.quantize`). Currently, only ``"nvfp4"`` and ``"mxfp8"``
+            are supported. Default: ``"nvfp4"``.
+    """
+
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        group_size: int = None,
+        bits: int = None,
+        mode: str = "nvfp4",
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
+        self.mode = mode
+
+        scale = math.sqrt(1 / input_dims)
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(output_dims, input_dims),
+        )
+        self._quantized = False
+
+    def _extra_repr(self):
+        out_dims, in_dims = self.weight.shape
+        if self.weight.dtype == mx.uint32:
+            in_dims = (in_dims * 32) // self.bits
+        return (
+            f"input_dims={in_dims}, output_dims={out_dims}, "
+            f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
+        )
+
+    def quantize(self):
+        if not self._quantized:
+            self.weight, self.scales = mx.quantize(
+                self.weight,
+                self.group_size,
+                self.bits,
+                mode=self.mode,
+            )
+            self._quantized = True
+
+    def dequantize(self):
+        if self._quantized:
+            self.weight = mx.dequantize(
+                self.weight,
+                scales=self.scales,
+                group_size=self.group_size,
+                bits=self.bits,
+                mode=self.mode,
+            )
+            self.__delattr__("scales")
+            self._quantized = False
+
+    def _set_training_mode(self, mode: bool):
+        super()._set_training_mode(mode)
+
+        if self._training:
+            self.dequantize()
+        else:
+            self.quantize()
+
+    def __call__(self, x):
+        x = mx.qqmm(
+            x,
+            self["weight"],
+            scales=self.get("scales"),
+            group_size=self.group_size,
+            bits=self.bits,
+            mode=self.mode,
+        )
+        return x
+
+    @classmethod
+    def from_linear(
+        cls,
+        linear_layer: Module,
+        group_size: int = None,
+        bits: int = None,
+        mode: str = "nvfp4",
+    ):
+        """Create a :obj:`QQLinear` layer from a :obj:`Linear` layer."""
+        output_dims, input_dims = linear_layer.weight.shape  # (N,K)
+        if linear_layer.get("bias") is not None:
+            raise NotImplementedError("QQLinear does not support bias yet.")
+        ql = cls(input_dims, output_dims, group_size, bits, mode=mode)
+        ql.weight = linear_layer.weight
+        ql.train(linear_layer.training)
+
+        return ql

data/mlx/python/mlx/nn/layers/recurrent.py

@@ -0,0 +1,289 @@
+# Copyright © 2024 Apple Inc.
+
+import math
+from typing import Callable, Optional
+
+import mlx.core as mx
+from mlx.nn.layers.activations import tanh
+from mlx.nn.layers.base import Module
+
+
+class RNN(Module):
+    r"""An Elman recurrent layer.
+
+    The input is a sequence of shape ``NLD`` or ``LD`` where:
+
+    * ``N`` is the optional batch dimension
+    * ``L`` is the sequence length
+    * ``D`` is the input's feature dimension
+
+    Concretely, for each element along the sequence length axis, this
+    layer applies the function:
+
+    .. math::
+
+        h_{t + 1} = \text{tanh} (W_{ih}x_t + W_{hh}h_t + b)
+
+    The hidden state :math:`h` has shape ``NH`` or ``H``, depending on
+    whether the input is batched or not. Returns the hidden state at each
+    time step, of shape ``NLH`` or ``LH``.
+
+    Args:
+        input_size (int): Dimension of the input, ``D``.
+        hidden_size (int): Dimension of the hidden state, ``H``.
+        bias (bool, optional): Whether to use a bias. Default: ``True``.
+        nonlinearity (callable, optional): Non-linearity to use. If ``None``,
+            then func:`tanh` is used. Default: ``None``.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool = True,
+        nonlinearity: Optional[Callable] = None,
+    ):
+        super().__init__()
+
+        self.nonlinearity = nonlinearity or tanh
+        if not callable(self.nonlinearity):
+            raise ValueError(
+                f"Nonlinearity must be callable. Current value: {nonlinearity}."
+            )
+
+        scale = 1.0 / math.sqrt(hidden_size)
+        self.hidden_size = hidden_size
+        self.Wxh = mx.random.uniform(
+            low=-scale, high=scale, shape=(hidden_size, input_size)
+        )
+        self.Whh = mx.random.uniform(
+            low=-scale, high=scale, shape=(hidden_size, hidden_size)
+        )
+        self.bias = (
+            mx.random.uniform(low=-scale, high=scale, shape=(hidden_size,))
+            if bias
+            else None
+        )
+
+    def _extra_repr(self):
+        return (
+            f"input_dims={self.Wxh.shape[1]}, "
+            f"hidden_size={self.hidden_size}, "
+            f"nonlinearity={self.nonlinearity}, bias={self.bias is not None}"
+        )
+
+    def __call__(self, x, hidden=None):
+        if self.bias is not None:
+            x = mx.addmm(self.bias, x, self.Wxh.T)
+        else:
+            x = x @ self.Wxh.T
+
+        all_hidden = []
+        for idx in range(x.shape[-2]):
+            if hidden is not None:
+                hidden = mx.addmm(x[..., idx, :], hidden, self.Whh.T)
+            else:
+                hidden = x[..., idx, :]
+            hidden = self.nonlinearity(hidden)
+            all_hidden.append(hidden)
+
+        return mx.stack(all_hidden, axis=-2)
+
+
+class GRU(Module):
+    r"""A gated recurrent unit (GRU) RNN layer.
+
+    The input has shape ``NLD`` or ``LD`` where:
+
+    * ``N`` is the optional batch dimension
+    * ``L`` is the sequence length
+    * ``D`` is the input's feature dimension
+
+    Concretely, for each element of the sequence, this layer computes:
+
+    .. math::
+
+        \begin{aligned}
+        r_t &= \sigma (W_{xr}x_t + W_{hr}h_t + b_{r}) \\
+        z_t &= \sigma (W_{xz}x_t + W_{hz}h_t + b_{z}) \\
+        n_t &= \text{tanh}(W_{xn}x_t + b_{n} + r_t \odot (W_{hn}h_t + b_{hn})) \\
+        h_{t + 1} &= (1 - z_t) \odot n_t + z_t \odot h_t
+        \end{aligned}
+
+    The hidden state :math:`h` has shape ``NH`` or ``H`` depending on
+    whether the input is batched or not. Returns the hidden state at each
+    time step of shape ``NLH`` or ``LH``.
+
+    Args:
+        input_size (int): Dimension of the input, ``D``.
+        hidden_size (int): Dimension of the hidden state, ``H``.
+        bias (bool): Whether to use biases or not. Default: ``True``.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        scale = 1.0 / math.sqrt(hidden_size)
+        self.Wx = mx.random.uniform(
+            low=-scale, high=scale, shape=(3 * hidden_size, input_size)
+        )
+        self.Wh = mx.random.uniform(
+            low=-scale, high=scale, shape=(3 * hidden_size, hidden_size)
+        )
+        self.b = (
+            mx.random.uniform(low=-scale, high=scale, shape=(3 * hidden_size,))
+            if bias
+            else None
+        )
+        self.bhn = (
+            mx.random.uniform(low=-scale, high=scale, shape=(hidden_size,))
+            if bias
+            else None
+        )
+
+    def _extra_repr(self):
+        return (
+            f"input_dims={self.Wx.shape[1]}, "
+            f"hidden_size={self.hidden_size}, bias={self.b is not None}"
+        )
+
+    def __call__(self, x, hidden=None):
+        if self.b is not None:
+            x = mx.addmm(self.b, x, self.Wx.T)
+        else:
+            x = x @ self.Wx.T
+
+        x_rz = x[..., : -self.hidden_size]
+        x_n = x[..., -self.hidden_size :]
+
+        all_hidden = []
+
+        for idx in range(x.shape[-2]):
+            rz = x_rz[..., idx, :]
+            if hidden is not None:
+                h_proj = hidden @ self.Wh.T
+                h_proj_rz = h_proj[..., : -self.hidden_size]
+                h_proj_n = h_proj[..., -self.hidden_size :]
+
+                if self.bhn is not None:
+                    h_proj_n += self.bhn
+
+                rz = rz + h_proj_rz
+
+            rz = mx.sigmoid(rz)
+
+            r, z = mx.split(rz, 2, axis=-1)
+
+            n = x_n[..., idx, :]
+
+            if hidden is not None:
+                n = n + r * h_proj_n
+            n = mx.tanh(n)
+
+            if hidden is not None:
+                hidden = (1 - z) * n + z * hidden
+            else:
+                hidden = (1 - z) * n
+
+            all_hidden.append(hidden)
+
+        return mx.stack(all_hidden, axis=-2)
+
+
+class LSTM(Module):
+    r"""An LSTM recurrent layer.
+
+    The input has shape ``NLD`` or ``LD`` where:
+
+    * ``N`` is the optional batch dimension
+    * ``L`` is the sequence length
+    * ``D`` is the input's feature dimension
+
+    Concretely, for each element of the sequence, this layer computes:
+
+    .. math::
+        \begin{aligned}
+        i_t &= \sigma (W_{xi}x_t + W_{hi}h_t + b_{i}) \\
+        f_t &= \sigma (W_{xf}x_t + W_{hf}h_t + b_{f}) \\
+        g_t &= \text{tanh} (W_{xg}x_t + W_{hg}h_t + b_{g}) \\
+        o_t &= \sigma (W_{xo}x_t + W_{ho}h_t + b_{o}) \\
+        c_{t + 1} &= f_t \odot c_t + i_t \odot g_t \\
+        h_{t + 1} &= o_t \text{tanh}(c_{t + 1})
+        \end{aligned}
+
+    The hidden state :math:`h` and cell state :math:`c` have shape ``NH``
+    or ``H``, depending on whether the input is batched or not.
+
+    The layer returns two arrays, the hidden state and the cell state at
+    each time step, both of shape ``NLH`` or ``LH``.
+
+    Args:
+        input_size (int): Dimension of the input, ``D``.
+        hidden_size (int): Dimension of the hidden state, ``H``.
+        bias (bool): Whether to use biases or not. Default: ``True``.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        scale = 1.0 / math.sqrt(hidden_size)
+        self.Wx = mx.random.uniform(
+            low=-scale, high=scale, shape=(4 * hidden_size, input_size)
+        )
+        self.Wh = mx.random.uniform(
+            low=-scale, high=scale, shape=(4 * hidden_size, hidden_size)
+        )
+        self.bias = (
+            mx.random.uniform(low=-scale, high=scale, shape=(4 * hidden_size,))
+            if bias
+            else None
+        )
+
+    def _extra_repr(self):
+        return (
+            f"input_dims={self.Wx.shape[1]}, "
+            f"hidden_size={self.hidden_size}, bias={self.bias is not None}"
+        )
+
+    def __call__(self, x, hidden=None, cell=None):
+        if self.bias is not None:
+            x = mx.addmm(self.bias, x, self.Wx.T)
+        else:
+            x = x @ self.Wx.T
+
+        all_hidden = []
+        all_cell = []
+
+        for idx in range(x.shape[-2]):
+            ifgo = x[..., idx, :]
+            if hidden is not None:
+                ifgo = mx.addmm(ifgo, hidden, self.Wh.T)
+            i, f, g, o = mx.split(ifgo, 4, axis=-1)
+
+            i = mx.sigmoid(i)
+            f = mx.sigmoid(f)
+            g = mx.tanh(g)
+            o = mx.sigmoid(o)
+
+            if cell is not None:
+                cell = f * cell + i * g
+            else:
+                cell = i * g
+            hidden = o * mx.tanh(cell)
+
+            all_cell.append(cell)
+            all_hidden.append(hidden)
+
+        return mx.stack(all_hidden, axis=-2), mx.stack(all_cell, axis=-2)