Group-Wise Int8 & Int4 Quantization

The problem

A model's weights are stored as numbers with 16 bits each. Quantization stores each weight as a small integer instead, plus one floating-point scale per group of weights that says how big one integer step is. The model shrinks (4 bits instead of 16 is four times smaller), and generation gets faster, because it is limited by reading memory. The price is a rounding error.

For $b$ bits, symmetric absmax quantization of a group of weights works like this:

and each weight is stored as $q_i = \mathrm{round}(w_i / s)$, an integer between $-q_{\max}$ and $q_{\max}$. To use the weight again, dequantize: $s \cdot q_i$.

where $s$ is the scale: the largest weight of the group (in absolute value) is mapped to $q_{\max}$, the others fall in between.

Example. One group of 4 weights, $b = 4$ bits: $w = (0.4, -1.2, 0.05, 0.9)$.

1. $q_{\max} = 2^3 - 1 = 7$ and $s = 1.2 / 7 = 0.171$.
2. $w / s = (2.33, -7, 0.29, 5.25)$, rounded to $q = (2, -7, 0, 5)$: small integers.
3. Dequantized: $s \cdot q = (0.343, -1.2, 0, 0.857)$. The errors, $(0.057, 0, 0.05, 0.043)$, are never above $s/2 = 0.086$.

With $b = 8$: $q_{\max} = 127$, $s = 0.0094$, $q = (42, -127, 5, 95)$ and the dequantized weights are $(0.397, -1.2, 0.047, 0.898)$: more bits, smaller error.

Groups

A single large weight makes $s$ large and ruins the precision of all the others in its group (they all round to 0). So the weights are cut into groups of group_size consecutive values, each with its own scale: an outlier only damages its own group. Storage is bits per weight plus one scale per group.

Example. 7 billion weights, 4 bits, groups of 128, scales of 2 bytes: $7 \times 10^9 \cdot 4 / 8 = 3.5$ GB (gigabytes) for the integers, plus $7 \times 10^9 / 128 = 54{,}687{,}500$ scales of 2 bytes, $0.109$ GB: $3.6$ GB in total, against $14$ GB in 16 bits.

What to implement

For a weight array w of any shape whose size is divisible by group_size:

• quantize(w, bits, group_size): flatten w, cut it into consecutive groups of group_size values, and return (q, scales): q an np.int8 array of shape (n_groups, group_size), scales a float array of shape (n_groups,).
• dequantize(q, scales, shape): the reconstructed float array, reshaped to shape.
• quantized_bytes(n_weights, bits, group_size, scale_bytes=2): the storage size: bits per weight (rounded up to whole bytes in total) plus one scale per group.

Requirements

• A group of all zeros has scale 0: its weights must come back as 0, without dividing by zero.
• bits is 8 or 4 (4-bit values are still stored in np.int8 here; real kernels pack two per byte).
• No Python loops over weights or groups.

Hint

w.reshape(-1, group_size) makes one row per group, and np.abs(groups).max(axis=1) gives the largest weight of each. To divide each row by its own scale, use scales[:, None].