Propagate quantization mode in quantized layers (#3133)

2026-02-15 18:33:13 -08:00
parent 43f4a74826
commit e226af720e
2 changed files with 46 additions and 7 deletions
@@ -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(