Add regression coverage for mixed-precision (per-layer) quantized checkpoint loading

[fdagostino]

What

Test-only PR: adds MixedPrecisionQuantLoadTests, end-to-end regression coverage for loading checkpoints that mix quantization precisions via per-module overrides in config.json — the layout every mlx-community Gemma 4 QAT upload uses (gemma-4-12B-it-qat-4bit, gemma-4-E4B-it-qat-4bit: 4-bit attention/embeddings, 8-bit mlp.{gate,up,down}_proj, declared as 144/126 per-module entries interleaved with the global quantization dict).

Why

Reports such as ollama/ollama#16740 attribute empty/garbage generation from these checkpoints to loaders reading the 8-bit MLP payload as 4-bit. While investigating exactly that symptom on device, I verified that mlx-swift-lm handles this correctly today: QuantizationContainer decodes the interleaved dict, PerLayerQuantization resolves each module path, and loadWeights applies the per-module bits (the same semantics as Python mlx-lm's class_predicate). Verified on macOS against the real mlx-community/gemma-4-12B-it-qat-4bit: after load the module tree has exactly 144 × QuantizedLinear(bits: 8) for the MLP projections, 4-bit everywhere else, and greedy generation is correct.

However, nothing in Tests/ exercises this path end to end (only the config-decode unit tests in BaseConfigurationTests), so the behavior the whole Gemma 4 QAT family depends on was unpinned. This PR pins it.

The tests

Built on a tiny text-only gemma4_unified model whose own initial weights are quantized into a synthetic on-disk checkpoint with the QAT layout:

1. Per-module quantization overrides are applied when loading — decodes a QAT-style quantization dict (global 4-bit + per-module 8-bit MLP overrides), runs loadWeights, and asserts each MLP projection comes back as QuantizedLinear(bits: 8, groupSize: 64), everything else quantizable at 4-bit, with an 8-bit dequantization round-trip check.
2. Ignoring the per-module overrides fails loudly, not silently — loading the same checkpoint with only the global 4-bit settings must throw (packed-shape mismatch) rather than silently mis-decode the 8-bit payload into garbage logits — the failure mode described in the ollama issue must never regress into silence.

Validated to bite: replacing the per-layer lookup in loadWeights with the global default makes test 1 fail with mismatchedSize on layers.0.mlp.gate_proj ([128, 8] vs [128, 16]).

Related

Follows the Gemma 4 series #390 (KV-shared construction), #391 (video), #392 (audio encoder).

🤖 Generated with Claude Code