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

Propagate quantization mode in quantized layers (#3133)

Browse Source
This commit is contained in:
vskiwi
2026-02-16 05:33:13 +03:00
committed by GitHub
parent 43f4a74826
commit e226af720e
2 changed files with 46 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mlx/nn/layers/distributed.py
+24 -6
View File
@@ -371,6 +371,8 @@ class QuantizedAllToShardedLinear(Module):
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``.
mode (str, optional): The quantization method to use (see
:func:`~mlx.core.quantize`). Default: ``"affine"``.
group (mx.distributed.Group, optional): The sharding will happen across
this group. If not set then the global group is used. Default is
``None``.
@@ -383,6 +385,7 @@ class QuantizedAllToShardedLinear(Module):
bias: bool = True,
group_size: int = 64,
bits: int = 4,
mode: str = "affine",
group: Optional[mx.distributed.Group] = None,
):
super().__init__()
@@ -390,6 +393,7 @@ class QuantizedAllToShardedLinear(Module):
# Quantization config
self.group_size = group_size
self.bits = bits
self.mode = mode
# Initialize the quantized weight
scale = math.sqrt(1.0 / input_dims)
@@ -406,7 +410,10 @@ class QuantizedAllToShardedLinear(Module):
high=scale,
shape=(output_dims // N, input_dims),
)
self.weight, self.scales, self.biases = mx.quantize(weight, group_size, bits)
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:
@@ -427,7 +434,7 @@ class QuantizedAllToShardedLinear(Module):
out_dims *= self.group.size()
return (
f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, "
f"group_size={self.group_size}, bits={self.bits}"
f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
)
def __call__(self, x: mx.array) -> mx.array:
@@ -438,10 +445,11 @@ class QuantizedAllToShardedLinear(Module):
x,
self["weight"],
scales=self["scales"],
biases=self["biases"],
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"]
@@ -465,6 +473,7 @@ class QuantizedAllToShardedLinear(Module):
hasattr(quantized_linear_layer, "bias"),
group_size=quantized_linear_layer.group_size,
bits=quantized_linear_layer.bits,
mode=getattr(quantized_linear_layer, "mode", "affine"),
group=group,
)
sl.update(
@@ -497,6 +506,8 @@ class QuantizedShardedToAllLinear(Module):
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``.
mode (str, optional): The quantization method to use (see
:func:`~mlx.core.quantize`). Default: ``"affine"``.
group (mx.distributed.Group, optional): The sharding will happen across
this group. If not set then the global group is used. Default is
``None``.
@@ -509,6 +520,7 @@ class QuantizedShardedToAllLinear(Module):
bias: bool = True,
group_size: int = 64,
bits: int = 4,
mode: str = "affine",
group: Optional[mx.distributed.Group] = None,
):
super().__init__()
@@ -516,6 +528,7 @@ class QuantizedShardedToAllLinear(Module):
# Quantization config
self.group_size = group_size
self.bits = bits
self.mode = mode
# Initialize the quantized weight
scale = math.sqrt(1.0 / input_dims)
@@ -532,7 +545,10 @@ class QuantizedShardedToAllLinear(Module):
high=scale,
shape=(output_dims, input_dims // N),
)
self.weight, self.scales, self.biases = mx.quantize(weight, group_size, bits)
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:
@@ -552,7 +568,7 @@ class QuantizedShardedToAllLinear(Module):
in_dims = (in_dims * 32) // self.bits * self.group.size()
return (
f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, "
f"group_size={self.group_size}, bits={self.bits}"
f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
)
def __call__(self, x: mx.array) -> mx.array:
@@ -560,10 +576,11 @@ class QuantizedShardedToAllLinear(Module):
x,
self["weight"],
scales=self["scales"],
biases=self["biases"],
biases=self.get("biases"),
transpose=True,
group_size=self.group_size,
bits=self.bits,
mode=self.mode,
)
x = mx.distributed.all_sum(x, group=self.group)
if "bias" in self:
@@ -588,6 +605,7 @@ class QuantizedShardedToAllLinear(Module):
hasattr(quantized_linear_layer, "bias"),
group_size=quantized_linear_layer.group_size,
bits=quantized_linear_layer.bits,
mode=getattr(quantized_linear_layer, "mode", "affine"),
group=group,
)
sl.update(
python/tests/mlx_distributed_tests.py
+22 -1
View File
@@ -146,7 +146,7 @@ class MLXDistributedCommonTestCase(mlx_tests.MLXTestCase):
self.assertTrue(mx.allclose(y, y2, atol=self.atol, rtol=self.rtol))
self.assertTrue(mx.allclose(y[part], y1, atol=self.atol, rtol=self.rtol))
# And their quant versions (QuintizedMatmul is not supported on CUDA)
# And their quant versions (QuantizedMatmul is not supported on CUDA)
if not mx.cuda.is_available():
qlin = lin.to_quantized()
slin1 = shard_linear(qlin, "all-to-sharded")
@@ -157,6 +157,27 @@ class MLXDistributedCommonTestCase(mlx_tests.MLXTestCase):
self.assertTrue(mx.allclose(y, y2, atol=self.atol, rtol=self.rtol))
self.assertTrue(mx.allclose(y[part], y1))
# Test non-affine quantization modes (mxfp8)
qlin_mxfp8 = lin.to_quantized(group_size=32, bits=8, mode="mxfp8")
self.assertEqual(qlin_mxfp8.mode, "mxfp8")
slin1_mxfp8 = shard_linear(qlin_mxfp8, "all-to-sharded")
slin2_mxfp8 = shard_linear(qlin_mxfp8, "sharded-to-all")
# Verify mode is propagated
self.assertEqual(slin1_mxfp8.mode, "mxfp8")
self.assertEqual(slin2_mxfp8.mode, "mxfp8")
# Verify biases parameter is not set for mxfp8
self.assertIsNone(slin1_mxfp8.get("biases"))
self.assertIsNone(slin2_mxfp8.get("biases"))
y = qlin_mxfp8(x)
y1 = slin1_mxfp8(x)
y2 = slin2_mxfp8(x[part])
self.assertTrue(mx.allclose(y, y2, atol=self.atol, rtol=self.rtol))
self.assertTrue(mx.allclose(y[part], y1))
# Check the backward works as expected
def dummy_loss(model, x, y):
return (model(x) * y).sum()