From e226af720e154c5fa90ee0938ab3d92f24b2df48 Mon Sep 17 00:00:00 2001 From: vskiwi <141816715+vskiwi@users.noreply.github.com> Date: Mon, 16 Feb 2026 05:33:13 +0300 Subject: [PATCH] Propagate quantization mode in quantized layers (#3133) --- python/mlx/nn/layers/distributed.py | 30 +++++++++++++++++++++------ python/tests/mlx_distributed_tests.py | 23 +++++++++++++++++++- 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/python/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index c7d79e55..86040479 100644 --- a/python/mlx/nn/layers/distributed.py +++ b/python/mlx/nn/layers/distributed.py @@ -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( diff --git a/python/tests/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index 8db45666..644da793 100644 --- a/python/tests/mlx_distributed_tests.py +++ b/python/tests/mlx_distributed_tests.py @@ -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()