quantize-py

python bindings for quantize, a simple, fast quantization library written in rust. quantization stores numbers in fewer bits, trading a little accuracy for a lot less memory.

pip install quantize-py

requires python 3.12 or newer, or 3.14 or newer for free-threaded python. numpy is installed with it. where no wheel fits, like alpine linux, pip builds it from source, which needs rust 1.88 or newer.

from quantize import Scale, quantize

weights = [0.42, -0.10, 0.70, -0.50]

q = quantize(weights, bits=8, block=32, scale=Scale.F16)
back = q.dequantize()  # [0.421, -0.098, 0.700, -0.498]
dot = q.dot(weights)  # 0.926

bits is the width of each code, from 2 to 16. block is how many values share one scale, and scale is how that scale is stored: Scale.F32 (the default), Scale.F16, or Scale.BF16. values can be a list, a numpy array, or anything else np.asarray reads, like a pytorch tensor. a 2-d array keeps its shape, so q.dequantize() gives back a matrix and q.matmul(x) computes x @ W.T, like a linear layer. both can write into a float32 numpy array or pytorch tensor you pass, like q.matmul(x, out=y), so a loop can reuse it. q.shape gives its rows and columns, and len(q) the number of values, as in rust.

the scales count toward the size: 4-bit codes with one f16 scale per 32 values cost 4.5 bits per value, or 5 with the default f32 scale. q.bits_per_element reports it. above about 10 bits, use f32 scales with asymmetric.quantize, since f16 and bf16 zero-points cap its accuracy.

the other schemes return the same Quantized type:

• asymmetric.quantize(weights, bits=8, block=32) adds a zero-point per block, for values that aren't centered on zero
• adaptive.quantize(weights, tolerance=0.1 * weights.std()) gives each block the fewest bits, from 2 to 8, that round every weight within tolerance, in the weights' own units. a tenth of their standard deviation gives about 5 bits a block. for a list, use np.std(weights)
• learned.refine(q, weights) refits each block's scale, and its zero-point if it has one, to lower the mean squared error. it changes q in place, so call q.copy() first to keep the original. it lets other threads run while it works, so threads can refit several layers at once
• learned.alternate(q, weights) refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes q in place and lets other threads run. both can raise the worst error past an adaptive tensor's tolerance, and lowering the error of the weights doesn't always lower the error of a model's outputs, so check those too
• Scheme.Q4_32.quantize(weights) picks a scheme at run time, and Scheme("symmetric(bits=4)") reads one from text, like a config value, which str(scheme) writes

a block with outliers can need more than 8 bits, which raises ToleranceTooTightError. retrying with its smallest_tolerance works, but loosens every block, not just that one:

import numpy as np
from quantize import ToleranceTooTightError, adaptive

weights = np.array(weights)  # a list has no .std(). skip this line for a pytorch tensor
try:
    q = adaptive.quantize(weights, tolerance=0.1 * weights.std())
except ToleranceTooTightError as error:
    q = adaptive.quantize(weights, tolerance=error.smallest_tolerance)

quantized values can be pickled, and compared with ==. q.to_bytes() saves one as bytes, in the same format as the rust crate, and Quantized.from_bytes(data) loads it back. to keep it in an np.savez or safetensors file, store np.frombuffer(q.to_bytes(), np.uint8).

to save its parts as plain arrays instead, like with np.savez, pass them back by name to Quantized.from_parts. leave out the one that's None: bits for an adaptive tensor, or block_bits for the others:

parts = dict(kind=q.kind, shape=q.shape, block=q.block, bits=q.bits, block_bits=q.block_bits,
             codes=q.codes, scales=q.scales, zero_points=q.zero_points, scale=q.scale.name)
np.savez("layer.npz", **{name: part for name, part in parts.items() if part is not None})
q = Quantized.from_parts(**np.load("layer.npz"))

to keep quantized values in a torch.save checkpoint, store torch.frombuffer(bytearray(q.to_bytes()), dtype=torch.uint8), and load each back with Quantized(t). the checkpoint is then as small as the bytes, and torch.load reads it without add_safe_globals. pickling q itself makes the checkpoint about 1.5 times larger, since torch.save stores bytes as text, and needs torch.serialization.add_safe_globals([Quantized]) before torch.load.

q.matmul runs on one core, but it lets other threads run while it multiplies, so threads can share out a batch. on an 8-core intel xeon, this multiplies a batch of 512 by a 4-bit 1536 × 576 matrix in 6 ms instead of 33 ms, with the same result, bit for bit:

import os
from concurrent.futures import ThreadPoolExecutor

pool = ThreadPoolExecutor()  # make it once, and reuse it for every layer

def linear(q, x):  # x has shape (batch, columns)
    # a piece per core, or pieces of 64 rows for a big batch, which stay in a core's cache
    pieces = np.array_split(x, max(os.cpu_count(), len(x) // 64))
    return np.concatenate(list(pool.map(q.matmul, pieces)))

out = linear(q, x)

with a gil, while another thread keeps running python, each piece can wait up to sys.getswitchinterval(), 5 ms by default, to get the gil back, so splitting pays off only while your other threads are idle or in native code, or on free-threaded python. dequantize, dot, and matmul on 65,536 values or fewer, like one vector times a 256 × 256 matrix, keep the gil, so they don't wait. quantizing and refitting keep it only up to 4,096 values, so threads can quantize a model's layers in parallel.

each value decodes as code * scale, or (code - zero_point) * scale with zero-points, using the scale and zero-point of its block. codes are signed and bits wide, and q.codes packs them low bits first. scales can be negative, since a symmetric block puts its value farthest from zero on the most negative code. help(Quantized) has the details.

to build and test from a clone of the repo, with rust 1.88 or newer and just:

just setup
just python

on debian or ubuntu, run sudo apt install python3-venv first.