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/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()