File size: 2,980 Bytes
23c283a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """Quantized weights for the port (MLX affine quantization), with a per-tensor recipe chosen by measurement.
A quantized weight is stored as `<name>.q` (packed uint32), `<name>.scales`, `<name>.biases`; `quant.json` in the
model folder records bits and group size per tensor. The runtime's linear layers call `mm`, which uses
`mx.quantized_matmul` for a quantized weight and a plain matmul otherwise; embedding tables use `rows`.
"""
import json
import os
import mlx.core as mx
class QW:
"""A quantized weight [out, in]: x @ W.T through mx.quantized_matmul."""
__slots__ = ("q", "scales", "biases", "bits", "group_size", "shape")
def __init__(self, q, scales, biases, bits, group_size, shape):
self.q, self.scales, self.biases = q, scales, biases
self.bits, self.group_size, self.shape = int(bits), int(group_size), tuple(shape)
@classmethod
def from_dense(cls, w, bits, group_size):
q, s, b = mx.quantize(w, group_size=group_size, bits=bits)
return cls(q, s, b, bits, group_size, w.shape)
def dense(self, dtype=None):
w = mx.dequantize(self.q, self.scales, self.biases, group_size=self.group_size, bits=self.bits)
return w if dtype is None else w.astype(dtype)
def arrays(self):
return [self.q, self.scales, self.biases]
def mm(x, w):
"""x @ w.T for a dense [out, in] weight or a QW."""
if isinstance(w, QW):
return mx.quantized_matmul(x, w.q, w.scales, w.biases, transpose=True, group_size=w.group_size, bits=w.bits)
return x @ w.T
def rows(w, idx):
"""Embedding lookup w[idx] for a dense table or a QW (dequantizes only the gathered rows)."""
if isinstance(w, QW):
return mx.dequantize(w.q[idx], w.scales[idx], w.biases[idx], group_size=w.group_size, bits=w.bits)
return w[idx]
def cast(w, dtype):
"""A dense weight is cast; a QW keeps its packed values and casts its scales and biases."""
if isinstance(w, QW):
return QW(w.q, w.scales.astype(dtype), w.biases.astype(dtype), w.bits, w.group_size, w.shape)
return w.astype(dtype)
def arrays(tree):
"""Every mx.array in a nested dict/list/QW structure (for mx.eval)."""
if isinstance(tree, QW):
return tree.arrays()
if isinstance(tree, dict):
return [a for v in tree.values() for a in arrays(v)]
if isinstance(tree, (list, tuple)):
return [a for v in tree for a in arrays(v)]
return [tree] if isinstance(tree, mx.array) else []
def load_recipe(path):
p = os.path.join(path, "quant.json")
return json.load(open(p)) if os.path.exists(p) else None
def take(w, name, recipe):
"""Pop `name` from a loaded file dict as a dense array or a QW, as quant.json says."""
if recipe and name in recipe["tensors"]:
r = recipe["tensors"][name]
return QW(w.pop(name + ".q"), w.pop(name + ".scales"), w.pop(name + ".biases"), r["bits"], r["group_size"],
r["shape"])
return w.pop(name)
|