diff --git a/python/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index f33e3cbc..c7d79e55 100644 --- a/python/mlx/nn/layers/distributed.py +++ b/python/mlx/nn/layers/distributed.py @@ -423,7 +423,7 @@ class QuantizedAllToShardedLinear(Module): def _extra_repr(self) -> str: out_dims, in_dims = self.weight.shape - in_dims *= 32 // self.bits + in_dims = (in_dims * 32) // self.bits out_dims *= self.group.size() return ( f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, " @@ -457,7 +457,7 @@ class QuantizedAllToShardedLinear(Module): ): group = group or mx.distributed.init() output_dims, input_dims = quantized_linear_layer.weight.shape - input_dims *= 32 // quantized_linear_layer.bits + input_dims = (input_dims * 32) // quantized_linear_layer.bits sl = cls( input_dims, @@ -549,7 +549,7 @@ class QuantizedShardedToAllLinear(Module): def _extra_repr(self) -> str: out_dims, in_dims = self.weight.shape - in_dims *= (32 // self.bits) * self.group.size() + 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}" @@ -580,7 +580,7 @@ class QuantizedShardedToAllLinear(Module): ): group = group or mx.distributed.init() output_dims, input_dims = quantized_linear_layer.weight.shape - input_dims *= 32 // quantized_linear_layer.bits + input_dims = (input_dims * 32) // quantized_linear_layer.bits sl = cls( input_dims, diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py index 1c98706d..0e520d15 100644 --- a/python/mlx/nn/layers/quantized.py +++ b/python/mlx/nn/layers/quantized.py @@ -233,7 +233,7 @@ class QuantizedLinear(Module): def _extra_repr(self): out_dims, in_dims = self.weight.shape - in_dims *= 32 // self.bits + in_dims = (in_dims * 32) // self.bits return ( f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, " f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}" @@ -338,7 +338,7 @@ class QQLinear(Module): def _extra_repr(self): out_dims, in_dims = self.weight.shape if self.weight.dtype == mx.uint32: - in_dims *= 32 // self.bits + in_dims = (in_dims * 32) // self.bits return ( f"input_dims={in_dims}, output_dims={out_dims}, " f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}" diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 6ded3722..b259da3e 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -211,6 +211,18 @@ class TestBase(mlx_tests.MLXTestCase): size = tree_reduce(lambda acc, p: acc + p.size, qlin.trainable_parameters(), 0) self.assertTrue(size > 0) + def test_quantized_sharded_linear_construction(self): + input_dims, output_dims = 1536, 1024 + for bits in [2, 3, 4, 5, 6, 8]: + lin = nn.Linear(input_dims, output_dims) + qlin = lin.to_quantized(bits=bits) + + slin1 = nn.QuantizedAllToShardedLinear.from_quantized_linear(qlin) + self.assertEqual(slin1.weight.shape, qlin.weight.shape) + + slin2 = nn.QuantizedShardedToAllLinear.from_quantized_linear(qlin) + self.assertEqual(slin2.weight.shape, qlin.weight.shape) + def test_grad_of_module(self): class Model(nn.Module): def __init__(self):