veloxquant_mlx.weight.model_quantizer module¶
Model-wide Linear-layer replacement with TurboQuant weight compression.
Provides quantize_model, which walks an mlx.nn.Module tree and
replaces each nn.Linear (and already affine-quantized mlx-lm
QuantizedLinear, which it first dequantizes) with a TurboQuant
QuantizedLinear (see weight/quantized_linear.py), skipping
embedding-like layers by name heuristic. compression_report then
summarizes the resulting memory savings across all replaced layers.
- veloxquant_mlx.weight.model_quantizer.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)