| import numpy as np |
| import pytest |
| from conftest import import_mlx |
|
|
| mx = import_mlx() |
| jax = pytest.importorskip("jax") |
|
|
| pytestmark = [pytest.mark.slow, pytest.mark.usefixtures("mlx_cpu")] |
|
|
|
|
| def _jax_logits(model, params, tokens): |
| from needle.model.quantize import WEIGHT_BITS, cq_ste_params |
|
|
| return model.apply({"params": cq_ste_params(params, WEIGHT_BITS)}, |
| jax.numpy.asarray(tokens), quant=True) |
|
|
|
|
| def _mlx_logits(config, params, tokens): |
| from needle.model.architecture_mlx import forward, to_mlx |
| from needle.model.quantize import WEIGHT_BITS |
| from needle.model.quantize_mlx import cq_ste_params |
|
|
| return forward(cq_ste_params(to_mlx(params), WEIGHT_BITS), config, mx.array(tokens), |
| quant=True) |
|
|
|
|
| @pytest.mark.parametrize("bits", [2, 4]) |
| def test_cq_quantize_matches_jax(bits): |
| from needle.model import quantize, quantize_mlx |
|
|
| w = np.random.default_rng(bits).standard_normal((3, 200, 96)).astype(np.float32) |
| ref = np.asarray(quantize.cq_quantize(jax.numpy.asarray(w), bits)) |
| got = np.asarray(quantize_mlx.cq_quantize(mx.array(w), bits)) |
| assert np.mean(np.abs(got - ref) > 1e-5) < 1e-3 |
| np.testing.assert_allclose(np.median(np.abs(got - ref)), 0, atol=1e-6) |
|
|
|
|
| @pytest.mark.parametrize("flash", [False, True]) |
| @pytest.mark.parametrize("kv_bits", [8, 4]) |
| def test_quantized_forward_matches_jax(flash, kv_bits, deploy_numerics, perturbed_model): |
| """An ulp-level difference can move a value across a rounding boundary; that shows |
| up as a few whole positions off, never as a drift across all of them.""" |
| model, config, params, tokens = perturbed_model(flash=flash, kv_bits=kv_bits) |
| deploy_numerics(config) |
| ref = np.asarray(_jax_logits(model, params, tokens)) |
| got = np.asarray(_mlx_logits(config, params, tokens)) |
| exact = np.all(np.abs(got - ref) <= 5e-5, axis=-1) |
| assert exact.mean() >= 0.9 |
| assert np.abs(got - ref).max() < 1e-2 |
|
|
|
|
| def test_published_checkpoint_quantized_loss_matches_jax(needle3_checkpoint, deploy_numerics): |
| """Under QAT the published model flips on ulp-level differences (JAX against itself |
| with 1e-7 weight noise agrees on ~96% of argmaxes), so it is compared by loss.""" |
| import optax |
| from needle.model.architecture import SimpleAttentionNetwork |
| from needle.model.run import load_checkpoint |
|
|
| params, config = load_checkpoint(needle3_checkpoint) |
| config.dtype = "float32" |
| params = jax.tree.map(lambda a: np.asarray(a, np.float32), params) |
| deploy_numerics(config) |
| tokens = np.random.default_rng(0).integers(14, config.vocab_size, (2, 96)).astype(np.int32) |
|
|
| def loss(logits): |
| return float(optax.softmax_cross_entropy_with_integer_labels( |
| np.asarray(logits)[:, :-1], tokens[:, 1:]).mean()) |
|
|
| ref = _jax_logits(SimpleAttentionNetwork(config), params, tokens) |
| got = _mlx_logits(config, params, tokens) |
| assert abs(loss(got) - loss(ref)) < 1e-2 * loss(ref) |
| assert (np.asarray(got).argmax(-1) == np.asarray(ref).argmax(-1)).mean() > 0.9 |
|
|
|
|
| def test_ab_scales_checkpoints_are_rejected(): |
| from needle.model.quantize_mlx import cq_ste_params |
|
|
| with pytest.raises(NotImplementedError): |
| cq_ste_params({"ab_scales/stack/a": mx.zeros((2, 2))}, 4) |
|
|
