| |
| |
| """A linear layer that runs the ternary kernel. |
| |
| ``TernaryLinear`` follows ``nn.Linear`` on the forward: ``y = x W^T + b``, |
| with ``W`` stored packed. It also replaces an ``nn.QuantizedLinear`` whose |
| affine 2-bit parameters are ternary (codes in {0, 1, 2}, bias = -scale). |
| """ |
| from __future__ import annotations |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
| import numpy as np |
|
|
| from .convert import from_mlx_2bit |
| from .kernel import tern_matmul |
| from .pack import GROUP, GROUPS, pack_trits |
|
|
|
|
| class TernaryLinear(nn.Module): |
| """Packed ternary affine layer. |
| |
| ``weight`` is uint8 ``[output_dims, input_dims / 4]``. ``scales`` is |
| float32 ``[output_dims, input_dims / group_size]``. ``group_size`` is 64 |
| (default) or 128. Parameters are frozen. |
| """ |
|
|
| def __init__(self, weight, scales, bias=None, group_size: int = GROUP): |
| super().__init__() |
| group_size = int(group_size) |
| if group_size not in GROUPS: |
| raise ValueError(f"group_size {group_size} not in {GROUPS}") |
| if not isinstance(weight, mx.array): |
| weight = mx.array(np.ascontiguousarray(weight)) |
| if weight.dtype != mx.uint8: |
| weight = weight.astype(mx.uint8) |
| if not isinstance(scales, mx.array): |
| scales = mx.array(np.asarray(scales, dtype=np.float32)) |
| if scales.dtype != mx.float32: |
| scales = scales.astype(mx.float32) |
| weight = mx.contiguous(weight) |
| scales = mx.contiguous(scales) |
| if weight.ndim != 2 or scales.ndim != 2: |
| raise ValueError("weight and scales must be matrices") |
| n, row_bytes = (int(v) for v in weight.shape) |
| k = row_bytes * 4 |
| if k % group_size or scales.shape != (n, k // group_size): |
| raise ValueError( |
| f"weight {tuple(weight.shape)} and scales {tuple(scales.shape)} " |
| f"are not a group-{group_size} ternary matrix" |
| ) |
| self.weight = weight |
| self.scales = scales |
| self.input_dims = k |
| self.output_dims = n |
| self.group_size = group_size |
| if bias is not None: |
| if not isinstance(bias, mx.array): |
| bias = mx.array(bias) |
| if tuple(bias.shape) != (n,): |
| raise ValueError(f"bias shape {tuple(bias.shape)} != ({n},)") |
| self.bias = bias |
| self.freeze() |
|
|
| def _extra_repr(self) -> str: |
| return ( |
| f"input_dims={self.input_dims}, output_dims={self.output_dims}, " |
| f"bias={'bias' in self}, group_size={self.group_size}, bits=2" |
| ) |
|
|
| def __call__(self, x): |
| if int(x.shape[-1]) != self.input_dims: |
| raise ValueError(f"x last dim {x.shape[-1]} != input_dims {self.input_dims}") |
| lead = x.shape[:-1] |
| flat = x.reshape(-1, self.input_dims) |
| y = tern_matmul(flat, self.weight, self.scales, self.input_dims, self.group_size) |
| if y.dtype != x.dtype: |
| y = y.astype(x.dtype) |
| if "bias" in self: |
| b = self.bias |
| if b.dtype != y.dtype: |
| b = b.astype(y.dtype) |
| y = y + b |
| return y.reshape(*lead, self.output_dims) |
|
|
| @classmethod |
| def from_trits(cls, trits, scales, bias=None, group_size: int = GROUP) -> "TernaryLinear": |
| """Build from int8 trits ``{-1,0,1}`` of shape ``[output, input]``.""" |
| packed, _row = pack_trits(np.asarray(trits)) |
| return cls(packed, scales, bias, group_size) |
|
|
| @classmethod |
| def from_mlx_2bit(cls, weight, scales, biases, bias=None, group_size: int = GROUP) -> "TernaryLinear": |
| """Build from MLX affine 2-bit ternary tensors.""" |
| packed, scale_f = from_mlx_2bit(weight, scales, biases, group_size) |
| return cls(packed, scale_f, bias, group_size) |
|
|
| @classmethod |
| def from_quantized_linear(cls, layer: nn.Module) -> "TernaryLinear": |
| """Replace one affine 2-bit ``QuantizedLinear`` that stores ternary weights.""" |
| if not isinstance(layer, nn.QuantizedLinear): |
| raise TypeError(f"expected QuantizedLinear, got {type(layer).__name__}") |
| if layer.bits != 2 or layer.group_size not in GROUPS or layer.mode != "affine": |
| raise ValueError( |
| f"expected affine 2-bit group {GROUPS}, got bits={layer.bits} " |
| f"group_size={layer.group_size} mode={layer.mode}" |
| ) |
| if layer.biases is None: |
| raise ValueError("affine 2-bit ternary storage needs biases equal to -scales") |
| extra = layer.bias if "bias" in layer else None |
| return cls.from_mlx_2bit( |
| layer.weight, layer.scales, layer.biases, bias=extra, group_size=layer.group_size |
| ) |
|
|
|
|
| def _replace_linear(value, path: str): |
| if isinstance(value, nn.QuantizedLinear): |
| try: |
| return TernaryLinear.from_quantized_linear(value), 1 |
| except ValueError as exc: |
| raise ValueError(f"{path}: {exc}") from exc |
| if isinstance(value, list): |
| n = 0 |
| for i, item in enumerate(list(value)): |
| new, k = _replace_linear(item, f"{path}.{i}") |
| if k: |
| value[i] = new |
| n += k |
| return value, n |
| if isinstance(value, dict) and not isinstance(value, nn.Module): |
| n = 0 |
| for key, item in list(value.items()): |
| new, k = _replace_linear(item, f"{path}.{key}") |
| if k: |
| value[key] = new |
| n += k |
| return value, n |
| return value, 0 |
|
|
|
|
| def replace_quantized_linears(model: nn.Module) -> int: |
| """Swap ternary affine 2-bit ``QuantizedLinear`` modules inside ``model``. |
| |
| The model object is updated in place. A ``QuantizedLinear`` that is not |
| ternary raises. Returns how many layers were replaced. |
| """ |
|
|
| replaced = 0 |
|
|
| def visit(path, mod): |
| nonlocal replaced |
| for key, value in list(mod.items()): |
| where = f"{path}.{key}" if path else str(key) |
| new, n = _replace_linear(value, where) |
| if n and new is not value: |
| mod[key] = new |
| replaced += n |
|
|
| model.apply_to_modules(visit) |
| return replaced |
|
|
