L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Metal/CPU nvfp4 and mxfp8 (#2946)

Browse Source
This commit is contained in:
Awni Hannun
2025-12-22 20:45:19 -08:00
committed by GitHub
parent 9cfda1a86e
commit 1eef1d155c
14 changed files with 660 additions and 536 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mlx/nn/layers/embedding.py
+7 -1
View File
@@ -1,6 +1,7 @@
# Copyright © 2023-2024 Apple Inc.
import math
from typing import Optional
import mlx.core as mx
from mlx.nn.layers.base import Module
@@ -39,6 +40,11 @@ class Embedding(Module):
"""
return x @ self.weight.T
def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"):
def to_quantized(
self,
group_size: Optional[int] = None,
bits: Optional[int] = None,
mode: str = "affine",
):
"""Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer."""
return QuantizedEmbedding.from_embedding(self, group_size, bits, mode)
python/mlx/nn/layers/linear.py
+7 -2
View File
@@ -1,7 +1,7 @@
# Copyright © 2023 Apple Inc.
import math
from typing import Any
from typing import Any, Optional
import mlx.core as mx
from mlx.nn.layers.base import Module
@@ -70,7 +70,12 @@ class Linear(Module):
x = x @ self["weight"].T
return x
def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"):
def to_quantized(
self,
group_size: Optional[int] = None,
bits: Optional[int] = None,
mode: str = "affine",
):
"""Return a :obj:`QuantizedLinear` layer that approximates this layer."""
return QuantizedLinear.from_linear(self, group_size, bits, mode)
python/mlx/nn/layers/quantized.py
+35 -26
View File
@@ -8,10 +8,21 @@ 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 = 64,
bits: int = 4,
group_size: int = None,
bits: int = None,
*,
mode: str = "affine",
class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = None,
@@ -24,10 +35,10 @@ def quantize(
Args:
model (mlx.nn.Module): The model whose leaf modules may be quantized.
group_size (int): The quantization group size (see
:func:`mlx.core.quantize`). Default: ``64``.
bits (int): The number of bits per parameter (see
:func:`mlx.core.quantize`). Default: ``4``.
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"``.
class_predicate (Optional[Callable]): A callable which receives the
@@ -72,10 +83,10 @@ class QuantizedEmbedding(Module):
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 (int, optional): The group size to use for the quantized
weight. See :func:`~mlx.core.quantize`. Default: ``64``.
bits (int, optional): The bit width to use for the quantized weight.
See :func:`~mlx.core.quantize`. Default: ``4``.
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"``.
"""
@@ -84,15 +95,14 @@ class QuantizedEmbedding(Module):
self,
num_embeddings: int,
dims: int,
group_size: int = 64,
bits: int = 4,
group_size: int = None,
bits: int = None,
mode: str = "affine",
):
super().__init__()
# Quantization config
self.group_size = group_size
self.bits = bits
self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
self.mode = mode
# Initialize the quantized weight
@@ -147,8 +157,8 @@ class QuantizedEmbedding(Module):
def from_embedding(
cls,
embedding_layer: Module,
group_size: int = 64,
bits: int = 4,
group_size: int = None,
bits: int = None,
mode: str = "affine",
):
"""Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer."""
@@ -179,10 +189,10 @@ class QuantizedLinear(Module):
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 (int, optional): The group size to use for the quantized
weight. See :func:`~mlx.core.quantize`. Default: ``64``.
bits (int, optional): The bit width to use for the quantized weight.
See :func:`~mlx.core.quantize`. Default: ``4``.
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"``.
"""
@@ -192,15 +202,14 @@ class QuantizedLinear(Module):
input_dims: int,
output_dims: int,
bias: bool = True,
group_size: int = 64,
bits: int = 4,
group_size: int = None,
bits: int = None,
mode: str = "affine",
):
super().__init__()
# Quantization config
self.group_size = group_size
self.bits = bits
self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
self.mode = mode
# Initialize the quantized weight
@@ -249,8 +258,8 @@ class QuantizedLinear(Module):
def from_linear(
cls,
linear_layer: Module,
group_size: int = 64,
bits: int = 4,
group_size: int = None,
bits: int = None,
mode: str = "affine",
):
"""Create a :obj:`QuantizedLinear` layer from a :obj:`Linear` layer."""