Tutorial 10: Quantization — Fitting Large Models in Memory¶
A model like Qwen3-8B stores its weights in float16 or bfloat16 — 2 bytes per parameter. At 8 billion parameters, that is 16 GB, which exceeds the memory of most Apple Silicon machines. Quantization compresses weights to 4 bits (or fewer), cutting memory usage by 4× while preserving most of the model's quality.
1. What Quantization Means¶
Neural network weights are floating-point numbers. Quantization replaces them with low-bit integers and stores a small number of scaling factors to recover approximate floating-point values:
Original: [0.12, -0.34, 0.56, -0.78, 0.90, -0.11, 0.23, -0.45]
↓ quantize to 4-bit
Packed: [0x32] ← 8 values packed into one 32-bit word
Scales: [0.15] ← one scale per group of 8 values
To recover the original values:
2. Group Quantization¶
The simplest quantization uses one scale per weight — terrible accuracy. Group
quantization uses one scale per group of G consecutive weights (typically
G=64):
Weights: [w0, w1, ..., w63, w64, w65, ..., w127]
Scale: [s0, s0, ..., s0, s1, s1, ..., s1 ]
← group 0 → ← group 1 →
Within each group, all weights share the same scale. The scale is computed as:
This maps the largest magnitude weight to the maximum integer value, and scales everything else proportionally.
3. The Math of 4-Bit Quantization¶
For 4-bit quantization (the most common for LLMs):
Encoding¶
import mlx.core as mx
def quantize_group(weights: mx.array, bits: int = 4) -> tuple:
"""
Quantize a group of weights to 4-bit integers.
Args:
weights: A group of weights. Shape: (group_size,)
bits: Number of bits per weight.
Returns:
(packed_ints, scale, bias) — the quantized representation.
"""
# Compute the scale: maps float range to integer range
qmax = (1 << (bits - 1)) - 1 # 7 for 4-bit (range: -8 to 7)
scale = mx.max(mx.abs(weights)) / qmax
# Quantize: round to nearest integer
integers = mx.round(weights / scale).astype(mx.int32)
# Clip to valid range
integers = mx.clip(integers, -(qmax + 1), qmax)
return integers, scale
Decoding¶
def dequantize_group(integers: mx.array, scale: mx.array, bits: int = 4) -> mx.array:
"""Convert 4-bit integers back to approximate floats."""
return integers.astype(mx.float32) * scale.astype(mx.float32)
Packing Multiple Weights into One Word¶
4 bits means each weight takes half a byte. We pack two weights per byte (or eight weights per 32-bit word):
def dequantize_weights(weight, scales, biases, group_size, bits):
"""
Dequantize packed low-bit weights to floating point.
Args:
weight: Packed integer weights. Shape: (..., packed_words)
scales: Per-group scales. Shape: (..., num_groups)
biases: Optional per-group biases. Shape: (..., num_groups) or None
group_size: Number of weights per group.
bits: Bits per weight (4, 8, etc.)
Returns:
Floating point weights. Shape: (..., num_groups * group_size)
"""
# How many low-bit values fit in one 32-bit word
values_per_word = 32 // bits # 8 values per word for 4-bit
# Create bit masks to extract each value from the packed word
shifts = mx.arange(0, 32, bits, dtype=mx.uint32)
mask = (1 << bits) - 1 # 0b1111 for 4-bit
# Extract individual values by shifting and masking
values = (weight[..., None] >> shifts) & mask
# Flatten from packed to unpacked representation
values = values.reshape(*weight.shape[:-1], weight.shape[-1] * values_per_word)
values = values.astype(mx.float32)
# Expand scales to match the unpacked shape (each scale covers group_size values)
expanded_scales = mx.repeat(scales, group_size, axis=-1).astype(mx.float32)
# Apply scales (and optional biases)
if biases is None:
return (values * expanded_scales).astype(scales.dtype)
expanded_biases = mx.repeat(biases, group_size, axis=-1).astype(mx.float32)
return (values * expanded_scales + expanded_biases).astype(scales.dtype)
4. The QuantizedWeights Container¶
Our engine wraps quantized data in a structured container:
class QuantizedWeights:
def __init__(self, scales, biases, group_size, bits, weight):
"""
Args:
weight: Packed low-bit integer weights.
scales: Per-group scale factors. Shape: (out_features, in_features // group_size)
biases: Optional per-group zero points. Shape: same as scales, or None.
group_size: Number of weights sharing one scale (typically 64).
bits: Bits per weight (typically 4).
"""
self.scales = scales # (out_features, num_groups)
self.biases = biases # (out_features, num_groups) or None
self.group_size = group_size
self.bits = bits
self.weight = weight # Packed int32 array
The memory savings are dramatic:
Original (bfloat16): 4096 × 14336 × 2 bytes = 112 MB
Quantized (4-bit): 4096 × 14336 × 0.5 bytes + scales ≈ 30 MB
≈ 3.7× compression
5. Quantized Matrix Multiply¶
The core operation is computing y = x @ W^T where W is quantized. There are
two approaches:
Approach 1: Dequantize Then Multiply¶
def dequantize_and_multiply(x, qw):
"""Simple path: dequantize weights, then do normal matmul."""
w_full = dequantize_weights(qw.weight, qw.scales, qw.biases,
qw.group_size, qw.bits)
return mx.matmul(x, w_full.T)
This is simple but wastes memory — the dequantized weight is 4× larger than the packed version.
Approach 2: Dequantize On-the-Fly (MLX Built-in)¶
def mlx_quantized_linear(x, w):
"""Use MLX's built-in quantized matmul — dequantizes in registers."""
return mx.quantized_matmul(
x, w.weight,
scales=w.scales,
biases=w.biases,
transpose=True,
group_size=w.group_size,
bits=w.bits,
)
MLX handles 4-bit dequantization inside the GPU kernel — the full float16 weight is never materialized in global memory.
Approach 3: Custom Metal Kernel¶
def quantized_linear(x, w):
"""Custom Metal kernel path — the fastest option when available."""
if native_available():
# Dispatch to Metal kernel with SIMD-optimized dequantization
return quantized_matmul(
w.scales, w.biases, w.group_size, w.bits,
x, w.weight, transpose_b=True,
)
# Fallback to MLX built-in
return mlx_quantized_linear(x, w)
The custom Metal kernel can use GPU-specific tricks like SIMD group operations to dequantize and multiply faster than the general-purpose MLX path.
6. Quantized Embeddings¶
Embeddings can also be quantized. Since embeddings are looked up by index (not multiplied), the dequantization is different:
class QuantizedEmbedding:
def __call__(self, x):
"""Look up rows from a quantized embedding matrix."""
# Select the relevant rows from packed weights, scales, and biases
return dequantize_weights(
self.weight.weight[x], # packed rows for these token IDs
self.weight.scales[x], # corresponding scales
None if self.weight.biases is None else self.weight.biases[x],
self.weight.group_size,
self.weight.bits,
)
7. When to Dequantize¶
Sometimes you need the full-precision version of a weight — for example, when debugging or when the quantized path is not available:
def dequantize_linear(weight):
"""Convert quantized weights back to bfloat16."""
if isinstance(weight, QuantizedWeights):
return mx.dequantize(
weight.weight,
weight.scales,
weight.biases,
weight.group_size,
weight.bits,
).astype(mx.bfloat16)
return weight # already full precision
8. Quantization Quality¶
4-bit quantization loses some quality, but it is surprisingly small:
| Model Size | FP16 PPL | 4-bit PPL | Quality Loss |
|---|---|---|---|
| 0.6B | 8.2 | 8.3 | ~1% |
| 4B | 7.1 | 7.2 | ~1.5% |
| 8B | 6.5 | 6.6 | ~1.5% |
PPL = perplexity (lower is better). The loss is typically less than 2% for models above 1B parameters. For local deployment, the memory savings far outweigh the quality loss.
9. Summary¶
- Quantization replaces float16 weights with 4-bit integers + per-group scales.
- Group quantization (group_size=64) balances compression and accuracy.
- 3 levels of quantized matmul: dequantize-then-multiply, MLX built-in, custom Metal kernel.
- 4-bit models use ~3.7× less memory with <2% quality loss.
- The quantized path is selected at load time based on the model's config.