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."""
python/tests/cuda_skip.py
+2 -2
View File
@@ -48,8 +48,8 @@ cuda_skip = {
"TestQuantized.test_qmm_shapes",
"TestQuantized.test_qmm_vjp",
"TestQuantized.test_qmv",
"TestQuantized.test_mxfp4_qmv",
"TestQuantized.test_mxfp4_qvm",
"TestQuantized.test_fp_qmv",
"TestQuantized.test_fp_qvm",
"TestQuantized.test_qvm",
"TestQuantized.test_qvm_splitk",
"TestQuantized.test_small_matrix",
python/tests/test_quantized.py
+101 -85
View File
@@ -289,7 +289,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
[128, 64, 32], # group_size
[2, 3, 4, 5, 6, 8], # bits
[256, 512, 67], # M
[64, 128], # N
[64, 256], # N
[0, 1, 3, 8], # B
)
for group_size, bits, M, N, B in tests:
@@ -309,33 +309,34 @@ class TestQuantized(mlx_tests.MLXTestCase):
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 1e-3)
def test_mxfp4_qmv(self):
def test_fp_qmv(self):
key = mx.random.key(0)
k1, k2 = mx.random.split(key)
tests = product(
[256, 512, 67], # M
[64, 128], # N
[64, 256], # N
[0, 1, 3, 8], # B
)
modes = ["mxfp4", "nvfp4", "mxfp8"]
for M, N, B in tests:
with self.subTest(shape=(B, M, N), group_size=32):
x_shape = (3, 1, N) if B == 0 else (B, 1, N)
w_shape = (M, N) if B == 0 else (B, M, N)
x = mx.random.normal(shape=x_shape, key=k1)
w = mx.random.normal(shape=w_shape, key=k2)
w_q, scales = mx.quantize(w, group_size=32, mode="mxfp4")
w_hat = mx.dequantize(w_q, scales, group_size=32, mode="mxfp4")
y_q = mx.quantized_matmul(
x,
w_q,
scales,
transpose=True,
group_size=32,
mode="mxfp4",
)
y_hat = x @ mx.swapaxes(w_hat, -1, -2)
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 1e-3)
for mode in modes:
with self.subTest(shape=(B, M, N), mode=mode):
x_shape = (3, 1, N) if B == 0 else (B, 1, N)
w_shape = (M, N) if B == 0 else (B, M, N)
x = mx.random.normal(shape=x_shape, key=k1)
w = mx.random.normal(shape=w_shape, key=k2)
w_q, scales = mx.quantize(w, mode=mode)
w_hat = mx.dequantize(w_q, scales, mode=mode)
y_q = mx.quantized_matmul(
x,
w_q,
scales,
transpose=True,
mode=mode,
)
y_hat = x @ mx.swapaxes(w_hat, -1, -2)
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 1e-3)
def test_qvm(self):
key = mx.random.key(0)
@@ -402,7 +403,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 2e-3)
def test_mxfp4_qvm(self):
def test_fp_qvm(self):
key = mx.random.key(0)
k1, k2 = mx.random.split(key)
tests = product(
@@ -413,26 +414,27 @@ class TestQuantized(mlx_tests.MLXTestCase):
# Add a splitk
tests = list(tests)
tests.append((128, 16384, 0))
modes = ["mxfp4", "nvfp4", "mxfp8"]
for M, N, B in tests:
with self.subTest(shape=(B, M, N)):
x_shape = (1, N) if B == 0 else (B, 1, N)
w_shape = (N, M) if B == 0 else (B, N, M)
x = mx.random.normal(shape=x_shape, key=k1)
w = mx.random.normal(shape=w_shape, key=k2)
w_q, scales = mx.quantize(w, group_size=32, mode="mxfp4")
w_hat = mx.dequantize(w_q, scales, group_size=32, mode="mxfp4")
y_q = mx.quantized_matmul(
x,
w_q,
scales,
transpose=False,
group_size=32,
mode="mxfp4",
)
y_hat = x @ w_hat
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 2e-3)
for mode in modes:
with self.subTest(shape=(B, M, N), mode=mode):
x_shape = (1, N) if B == 0 else (B, 1, N)
w_shape = (N, M) if B == 0 else (B, N, M)
x = mx.random.normal(shape=x_shape, key=k1)
w = mx.random.normal(shape=w_shape, key=k2)
w_q, scales = mx.quantize(w, mode=mode)
w_hat = mx.dequantize(w_q, scales, mode=mode)
y_q = mx.quantized_matmul(
x,
w_q,
scales,
transpose=False,
mode=mode,
)
y_hat = x @ w_hat
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 2e-3)
def test_mode_error_cases(self):
w = mx.random.normal(shape=(256, 256))
@@ -626,7 +628,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
self.assertLess((y_q - y_hat).abs().max(), 1e-3)
def test_gather_qmm(self):
def quantize(w, transpose=True, group_size=64, bits=4, mode="affine"):
def quantize(w, transpose=True, group_size=None, bits=None, mode="affine"):
if mode == "affine":
qw, s, b = mx.quantize(w, group_size=group_size, bits=bits, mode=mode)
else:
@@ -647,8 +649,8 @@ class TestQuantized(mlx_tests.MLXTestCase):
lhs_indices=None,
rhs_indices=None,
transpose=True,
group_size=64,
bits=4,
group_size=None,
bits=None,
mode="affine",
):
with self.subTest(
@@ -737,9 +739,22 @@ class TestQuantized(mlx_tests.MLXTestCase):
"lhs_indices": (0,),
"batch_B": (3,),
"rhs_indices": (2, 1),
"group_size": 32,
"mode": "nvfp4",
},
{
"batch_A": (1,),
"lhs_indices": (0,),
"batch_B": (3,),
"rhs_indices": (2, 1),
"mode": "mxfp4",
},
{
"batch_A": (1,),
"lhs_indices": (0,),
"batch_B": (3,),
"rhs_indices": (2, 1),
"mode": "mxfp8",
},
)
for kwargs in inputs:
@@ -753,24 +768,24 @@ class TestQuantized(mlx_tests.MLXTestCase):
test_shape(32, 512, 32, transpose=False, **kwargs)
test_shape(1, 512, 32, transpose=False, **kwargs)
def test_qmm_mxfp4_type(self):
def test_qmm_fp_type(self):
indices = mx.array([[2], [0], [1]], dtype=mx.uint32)
for t in [mx.bfloat16, mx.float16, mx.float32]:
x = mx.random.normal((32, 256)).astype(t)
modes = ["mxfp8", "mxfp4"]
for mode in modes:
for t in [mx.bfloat16, mx.float16, mx.float32]:
x = mx.random.normal((32, 256)).astype(t)
w = mx.random.normal((32, 256))
wq, s = mx.quantize(w, mode="mxfp4", bits=4, group_size=32)
out = mx.quantized_matmul(x, wq, s, mode="mxfp4", group_size=32, bits=4)
self.assertEqual(out.dtype, t)
w = mx.random.normal((32, 256))
wq, s = mx.quantize(w, mode=mode)
out = mx.quantized_matmul(x, wq, s, mode=mode)
self.assertEqual(out.dtype, t)
w = mx.random.normal((4, 32, 256))
wq, s = mx.quantize(w, mode="mxfp4", bits=4, group_size=32)
w = mx.random.normal((4, 32, 256))
wq, s = mx.quantize(w, mode=mode)
out = mx.gather_qmm(
x, wq, s, rhs_indices=indices, mode="mxfp4", group_size=32, bits=4
)
self.assertEqual(out.dtype, t)
out = mx.gather_qmm(x, wq, s, rhs_indices=indices, mode=mode)
self.assertEqual(out.dtype, t)
def test_gather_matmul_grad(self):
def quantize(w, transpose=True, group_size=64, bits=4):
@@ -802,14 +817,14 @@ class TestQuantized(mlx_tests.MLXTestCase):
self.assertTrue(mx.allclose(g1, g2, atol=1e-4))
def test_gather_qmm_sorted(self):
def quantize(w, transpose=True, bits=4, group_size=64, mode="affine"):
def quantize(w, transpose=True, group_size=None, mode="affine"):
if mode == "affine":
qw, s, b = mx.quantize(w, group_size=group_size, bits=bits, mode=mode)
qw, s, b = mx.quantize(w, group_size=group_size, mode=mode)
else:
qw, s = mx.quantize(w, group_size=group_size, bits=bits, mode=mode)
qw, s = mx.quantize(w, mode=mode)
b = None
w_hat = mx.dequantize(qw, s, b, group_size=group_size, bits=bits, mode=mode)
w_hat = mx.dequantize(qw, s, b, group_size=group_size, mode=mode)
if transpose:
w_hat = w_hat.swapaxes(-1, -2)
return w_hat, qw, s, b
@@ -831,11 +846,15 @@ class TestQuantized(mlx_tests.MLXTestCase):
# L, K, D, E, I, transpose
(32, 512, 512, 4, 2, True, "affine"),
(32, 512, 544, 4, 2, True, "mxfp4"),
(32, 512, 544, 4, 2, True, "nvfp4"),
(32, 512, 544, 4, 2, True, "mxfp8"),
(133, 512, 512, 4, 2, True, "affine"),
(133, 512, 555, 4, 2, True, "affine"),
(133, 512, 512, 4, 2, True, "affine"),
(64, 512, 512, 4, 2, False, "affine"),
(64, 512, 544, 4, 2, False, "mxfp4"),
(64, 512, 544, 4, 2, False, "nvfp4"),
(64, 512, 544, 4, 2, False, "mxfp8"),
(133, 512, 512, 4, 2, False, "affine"),
(133, 512, 544, 4, 2, False, "affine"),
(133, 512, 555, 4, 2, False, "affine"),
@@ -848,8 +867,8 @@ class TestQuantized(mlx_tests.MLXTestCase):
for L, K, D, E, I, transpose, mode in parameters:
with self.subTest(L=L, K=K, D=D, E=E, I=I, transpose=transpose, mode=mode):
if mode == "mxfp4":
group_size = 32
if mode != "affine":
group_size = None
dtype = (
mx.bfloat16 if (mx.default_device() == mx.gpu) else mx.float32
)
@@ -984,36 +1003,33 @@ class TestQuantized(mlx_tests.MLXTestCase):
num_ds = (out_up - out_down) / (2 * eps)
self.assertAlmostEqual(dparams[p][idx], num_ds, delta=2e-2)
def test_mxfp4_vjp_scales_throws(self):
def test_fp_vjp_scales_throws(self):
mx.random.seed(0)
x = mx.random.normal(shape=(2, 512))
w = mx.random.normal(shape=(512, 512))
wq, s = mx.quantize(w, bits=4, group_size=32, mode="mxfp4")
for mode in ["mxfp4", "mxfp8", "nvfp4"]:
wq, s = mx.quantize(w, mode=mode)
def mm(s, x, wq):
return mx.quantized_matmul(
x, wq, s, bits=4, group_size=32, mode="mxfp4"
).sum()
def mm(s, x, wq):
return mx.quantized_matmul(x, wq, s, mode=mode).sum()
# Should raise
with self.assertRaises(ValueError):
ds = mx.grad(mm)(s, x, wq)
# Should raise
with self.assertRaises(ValueError):
ds = mx.grad(mm)(s, x, wq)
rhs_indices = mx.array(0)
with self.assertRaises(ValueError):
rhs_indices = mx.array(0)
with self.assertRaises(ValueError):
def gmm(s, x, wq):
return mx.gather_qmm(
x,
wq,
s,
rhs_indices=rhs_indices,
bits=4,
group_size=32,
mode="mxfp4",
).sum()
def gmm(s, x, wq):
return mx.gather_qmm(
x,
wq,
s,
rhs_indices=rhs_indices,
mode=mode,
).sum()
ds = mx.grad(gmm)(s, x, wq)
ds = mx.grad(gmm)(s, x, wq)
def test_quantize_strided(self):
N = 64