veloxquant_mlx.weight package¶
Whole-model weight compression, built on the TurboQuant quantizer family.
Groups quantize_model/compression_report for replacing an
mlx.nn model’s Linear layers with QuantizedLinear (TurboQuant-style
rotate + Lloyd-Max weight compression), and
save_reservoir/load_reservoir/graft_reservoir for serializing
and restoring a quantized model’s compressed weights to/from a single flat
file.
- class veloxquant_mlx.weight.QuantizedLinear(in_features, out_features, bits=4, use_hadamard=True, bias=True, seed=42)[source]
Bases:
ModuleDrop-in replacement for nn.Linear with TurboQuant weight compression.
- Weight matrix W (out, in) is quantized row-by-row offline:
Normalize each row: w_norm = w / ||w||_2 (store norm)
Rotate: w_rot = preconditioner.apply(w_norm)
Quantize to b-bit Lloyd-Max indices (codebook fits N(0, 1/sqrt(in)))
Store indices + per-row norms
- Forward pass:
w_hat = norms * unrotate(centroids[indices]) out = x @ w_hat.T + bias
Per-row normalization is critical: weight rows have varying magnitudes but the Lloyd-Max codebook is calibrated for unit-norm post-rotation vectors. Without normalization the codebook range is completely wrong.
- Parameters:
in_features (
int) – Input dimension.out_features (
int) – Output dimension.bits (
int) – Bit-width for weight quantization (2, 3, or 4).use_hadamard (
bool) – Use randomized Hadamard (O(d log d)) instead of QR (O(d²)). Auto-falls back to QR if in_features is not Hadamard-compatible.bias (
bool) – Whether the layer has a bias term.seed (
int) – Random seed for the rotation matrix.
- quantize_weights(weight, bias=None)[source]
Compress a weight matrix into this layer.
- Parameters:
weight (mx.array) – Shape (out_features, in_features), fp16 or fp32.
bias (mx.array | None) – Optional shape (out_features,).
- Return type:
None
- property memory_bytes: int
indices + norms (excludes codebook, shared).
- Type:
Compressed storage
- property fp16_bytes: int
Uncompressed fp16 weight matrix size, in bytes.
- veloxquant_mlx.weight.compression_report(model)[source]
Summarise memory savings across all QuantizedLinear layers.
- Parameters:
model (
Module) – Model after quantize_model() has been called.- Return type:
- Returns:
Dict with total_compressed_bytes, total_fp16_bytes, ratio, n_layers.
- veloxquant_mlx.weight.quantize_model(model, bits=4, use_hadamard=True, skip_embeddings=True, seed=42)[source]
Replace all Linear layers in a model with TurboQuant QuantizedLinear.
Handles both: - nn.Linear (fp16 weights) — compressed directly - mlx-lm’s QuantizedLinear (affine-quantized) — dequantized first, then
re-quantized with TurboQuant’s rotation + Lloyd-Max scheme
- Parameters:
model (
Module) – Any mlx.nn.Module (e.g. a loaded mlx-lm model).bits (
int) – Bit-width for TurboQuant compression (2, 3, or 4).use_hadamard (
bool) – Use Metal-accelerated Hadamard rotation (recommended).skip_embeddings (
bool) – Skip embedding-like layers.seed (
int) – Base random seed. Each layer gets seed + layer_index.
- Return type:
Module- Returns:
The same model with Linear layers replaced (in-place mutation + return).
Example:
import mlx_lm model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit") model = quantize_model(model, bits=3, use_hadamard=True)
- veloxquant_mlx.weight.save_reservoir(model, path, persist_rotation=False)[source]
Serialize a quantize_model()-processed model’s QuantizedLinear layers.
Writes a single flat file: magic/version, a JSON header describing each layer’s shape and quantization parameters, followed by page-aligned index/norms/rotation/centroids blobs. See docs/WEIGHT_RESERVOIR_IDEATION.md for format rationale and the measured findings on why this is a load-time / compressed-copy-size optimization rather than a cross-process RAM-sharing one.
- Parameters:
model (
Module) – A model whose Linear layers have already been replaced by quantize_model() (or manually constructed QuantizedLinear layers with quantize_weights() called).persist_rotation (
bool) – If True, persist each QR-fallback layer’s full d x d rotation matrix so load_reservoir() never calls np.linalg.qr. This makes load ~9-230x faster (measured; see Finding 3 in docs/WEIGHT_RESERVOIR_IDEATION.md) but a dense random d x d orthogonal matrix is inherently O(d^2) information with no compact encoding, so this can make the file MUCH larger than the source model for architectures with non-Hadamard- compatible layer widths (measured: 9.5x larger for Qwen2.5-0.5B-Instruct-4bit, Finding 4). Default False: rotation matrices are re-derived from the layer’s seed on load (same QR cost quantize_model() already pays), keeping the file close to the size of the compressed weights themselves. Hadamard- compatible layers persist their (d,) diagonal either way – that storage is cheap regardless of this flag.
- Return type:
- veloxquant_mlx.weight.load_reservoir(path)[source]
Deserialize a reservoir file into a flat {layer_name: QuantizedLinear} dict.
Reconstructs each layer from its persisted rotation matrix/diagonal and codebook centroids (not by recomputing them from the seed – see _fast_quantized_linear), then assigns the persisted _w_indices / _w_norms directly. This skips both the dequantize -> rotate -> Lloyd-Max-argmin quantization pass AND the rotation/codebook setup cost, which for QR-fallback layers otherwise dominates load time.
The file is memory-mapped (np.memmap) for the read; per Finding 2 in docs/WEIGHT_RESERVOIR_IDEATION.md, mx.array() still performs one copy out of that mapping into MLX’s unified-memory arena — there is no zero-copy constructor in the installed MLX version. The win is a smaller copy (2-4 bit indices instead of fp16 weights) and no re-quantization/re-setup compute, not zero RAM cost.
- veloxquant_mlx.weight.graft_reservoir(model, path)[source]
Load a reservoir and replace the matching named modules in model in place.
model must have Linear/QuantizedLinear layers at the same dotted names the reservoir was saved with (i.e. it should be the same architecture quantize_model() would have walked). Layers not present in the reservoir are left untouched.