Fix sharding of quantized models with non-power-of-2 bits (#3006)

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,

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}"

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):