"""Tests for quantize.py: the hand-worked examples of the lesson, the claims it makes, and the sizes it quotes.""" import numpy as np import pytest import quantize as qz CFG = qz.load_config() def test_hand_worked_four_bit_example(): """Guards the scale or rounding rule drifting from the table in the lesson: absmax 0.70 / 7 = 0.1, then round.""" w = [0.32, -0.10, 0.04, -0.70] codes, scales = qz.quantize(w, 4, group=4) assert codes.tolist() == [3, -1, 0, -7] assert float(scales[0, 0]) == pytest.approx(0.1, rel=1e-3) # stored as a 16-bit float, so 0.0999756 assert qz.dequantize(codes, scales, group=4).round(2).tolist() == [0.3, -0.1, 0.0, -0.7] def test_hand_worked_three_bit_example(): """Guards the 3-bit row of the lesson: only seven codes (-3 to 3), so small weights fall to zero.""" codes, _ = qz.quantize([0.32, -0.10, 0.04, -0.70], 3, group=4) assert codes.tolist() == [1, 0, 0, -3] def test_ties_round_to_even(): """Guards round-half-up creeping in (Python's round and numpy's rint agree; JavaScript's Math.round does not).""" codes, scales = qz.quantize([3.0, 2.5, 0.5, -0.5, -1.5, -2.5], 3) assert float(scales[0, 0]) == 1.0 assert codes.tolist() == [3, 2, 0, 0, -2, -2] @pytest.mark.parametrize("bits", [8, 4, 3]) @pytest.mark.parametrize("group", [None, 128]) def test_codes_stay_in_range_and_the_largest_weight_survives(bits, group): """Guards a scale cut by qmax = 2**bits (too wide) or a 16-bit scale that rounds down and overflows the range.""" w = qz.synthetic_weights(8, 256, seed=3, outlier=0.5) codes, scales = qz.quantize(w, bits, group) assert np.abs(codes).max() == qz.qmax(bits) back = qz.dequantize(codes, scales, group) assert abs(back).max() == pytest.approx(0.5, rel=2e-3) def test_a_group_of_zeros_does_not_divide_by_zero(): """Guards the zero-step guard: 0 / 0 is NaN, which numpy casts to 0 on arm64 but to INT_MIN on x86, so only a raised error shows it.""" w = np.zeros((2, 8)) w[1, :4] = [0.1, -0.2, 0.3, 0.05] with np.errstate(invalid="raise", divide="raise"): codes, scales = qz.quantize(w, 4, group=4) back = qz.dequantize(codes, scales, 4) assert not np.isnan(back).any() assert codes[0].tolist() == [0] * 8 def test_group_must_divide_the_row(): with pytest.raises(ValueError): qz.quantize(np.ones((2, 10)), 4, group=4) def test_error_grows_as_bits_fall_and_groups_help(): w = qz.synthetic_weights() per_tensor = [qz.error(w, b) for b in (8, 4, 3)] per_group = [qz.error(w, b, qz.GROUP) for b in (8, 4, 3)] assert per_tensor == sorted(per_tensor) and per_group == sorted(per_group) assert all(g < t for g, t in zip(per_group, per_tensor)) assert qz.error(w, 16) == 0.0 def test_one_outlier_ruins_a_shared_scale_but_not_groups(): """Guards the lesson's central picture: 1 weight in 262,144 stretches the shared scale until most weights round to zero. Also pins the four errors the lesson quotes (19.5, 11.8, 99.5, 11.9), which move if the seed of the made-up matrix changes.""" clean, spiked = qz.synthetic_weights(), qz.synthetic_weights(outlier=1.0) assert [round(100 * qz.error(w, 4, g), 1) for w in (clean, spiked) for g in (None, qz.GROUP)] == [19.5, 11.8, 99.5, 11.9] assert qz.error(spiked, 4) > 0.9 assert qz.error(spiked, 4, qz.GROUP) - qz.error(clean, 4, qz.GROUP) < 0.005 assert np.count_nonzero(qz.roundtrip(spiked, 4)) < 0.01 * spiked.size def test_widget_sample_opens_on_the_numbers_the_lesson_quotes(): """Guards the sample weights or outlier slot drifting from the lesson's 4.7%, 19.1% (big weight) and 34.7% (big weight, one scale).""" assert max(map(abs, qz.sample_weights())) == 0.283 assert [round(100 * qz.sample_error(4, 4, False), 1), round(100 * qz.sample_error(4, 4, True), 1), round(100 * qz.sample_error(4, None, True), 1)] == [4.7, 19.1, 34.7] def test_bits_per_weight_counts_the_scale(): assert qz.bits_per_weight(4, 128) == 4.125 and qz.bits_per_weight(8, 128) == 8.125 and qz.bits_per_weight(3, 128) == 3.125 assert qz.bits_per_weight(16, 128) == 16 and qz.bits_per_weight(4, None) == 4 def test_llama_sizes_match_lesson_18_plus_scales(): """Guards the size formula against lesson 18's weights (8,030,261,248; 16 bits is 14.96 GiB) and the added scales.""" assert qz.param_count(CFG) == 8_030_261_248 assert qz.weight_bytes(CFG, 16) == 16_060_522_496 others = qz.param_count(CFG) - qz.embedding_params(CFG) assert qz.weight_bytes(CFG, 4, None) == qz.embedding_params(CFG) * 2 + others // 2 assert qz.weight_bytes(CFG, 4) - qz.weight_bytes(CFG, 4, None) == others // 128 * 2 assert [round(qz.weight_bytes(CFG, b) / qz.GIB, 2) for b in qz.BITS] == [14.96, 8.56, 5.31, 4.50] # Sizes of the `.safetensors` files on Hugging Face (Hub API, read 2026-10-19), Llama 3.1 8B Instruct. AWQ_INT4_FILE = 5_727_938_576 # AutoAWQ, 4 bits, group 128, scale and 4-bit zero point per group W8A16_FILE = 9_084_040_952 # llm-compressor, 8 bits, one scale per output row def test_formula_lands_on_published_files(): """Guards the size formula against what real files weigh. AWQ stores a 4-bit zero point beside each scale (half a byte).""" others = qz.param_count(CFG) - qz.embedding_params(CFG) awq = qz.weight_bytes(CFG, 4) + others // 128 // 2 assert abs(awq - AWQ_INT4_FILE) / AWQ_INT4_FILE < 1e-4 assert qz.weight_bytes(CFG, 8, None) < W8A16_FILE < qz.weight_bytes(CFG, 8, 128) def test_decode_reads_one_row_of_the_token_table_not_all_of_it(): """Guards counting the 1.05 GB token table as read per token when only one 8 KB row is.""" table = qz.table_params(CFG) * 2 assert qz.bytes_per_token(CFG, 16) == qz.weight_bytes(CFG, 16) - table + 4096 * 2 tied = dict(CFG, tie_word_embeddings=True) assert qz.bytes_per_token(tied, 16) == qz.weight_bytes(tied, 16) # a tied table is also the output matrix: read in full def test_speed_ceiling_rises_as_bits_fall(): h100 = qz.BANDWIDTHS["H100 SXM"] ceilings = [qz.tokens_per_second_bound(qz.bytes_per_token(CFG, b), h100) for b in qz.BITS] assert ceilings == sorted(ceilings) assert [round(c) for c in ceilings] == [223, 412, 721, 887]