Quantized inference in posttrainllm's HFModel — Swift-side QuantizedLinear

Quantized inference in posttrainllm’s HFModel — Swift-side QuantizedLinear

Status: PRD (2026-06-09). Phase 1 SHIPPED 2026-06-10 (serve --quantize); native packed load (the design below) still open.

2026-06-10 Phase 1 findings — serve --quantize int4|int8

Shipped (commit d3b7879): in-memory MLXNN.quantize after load, same mechanic as sample --quantize. Skipped under --lora; for adapters, bake first (bake-lora → serve --quantize). Quantize config folded into the KV prompt-cache fingerprint.

Measured on the shipping planner (pace-planner-v9-lora/baked-hf, Qwen3-0.6B, M5 Pro, release build, fm-fixtures-v2 with v9 compose-v2 prompt + v9 schema + serve-side --grammar):

• int8 is the free lunch: 2.3× decode, identical fixture score. int4 is faster still but drops a fixture — per the no-quality-regression rule, int8 is the planner default and int4 needs a per-model gate run before use.
• RSS is flat (~3 GB): load is still fp16-materialize-then-quantize. The real memory/disk win needs the native packed load below — that remains the open Phase 2.
• Two eval-setup traps (cost ~30 min): per-request response_format does NOT trigger constrained decoding — the grammar must be passed serve-side via --grammar, or the model free-runs <think> and scores 0%. And the frozen 33.3% Dim1 baseline appears to have used a different prompt/schema config: with v9 compose-v2 prompt + v9 schema + grammar, fp16 v9 scores 60.0% — worth re-examining the Dim1 baseline before the v11 run.

Problem

mlx_lm convert -q --q-bits 4 produces MLX-quantized safetensors:

• Weights stored as uint32 (8 values per uint32 for 4-bit)
• Companion tensors per quantized weight: <name>.scales, <name>.biases
• config.json has quantization block with bits, group_size, mode

Tinygpt’s HFModel.swift::makeMLXArray (line 442) fatal-errors on U32 dtype. The loader has no concept of quantized weights.

Result: v9-LoRA baked-hf is 1.4 GB fp16. Quantized to Q4 is 331 MB (4.2× smaller). But posttrainllm serve cannot load it → Pace ships fp16.

Target outcome

Tinygpt serve loads mlx_lm convert -q output natively. Pace bundles 331 MB instead of 1.4 GB. Inference uses MLX’s quantized_matmul ops directly on packed weights — no dequant-at-load.

Estimated formula payoff:

• cost: 4× lower disk + RAM
• speed: 1.2-1.5× faster on M-series (quantized matmul is bandwidth-bound; fewer bytes to move)
• accuracy: ~1-2% drop on Q4 (per published Qwen3 quant benchmarks)
• Net formula: ~30-40% boost — clears the 5%/2-week bar 6-8× over

Design

1. Detect quantization at load

In HFModelLoader.load(), parse config.json["quantization"] (or quantization_config). If present, set a quantized: QuantizationSpec on the load result with {bits, groupSize, mode}.

2. Detect quantized tensors during weight load

HFModel.swift::loadShard, walking each tensor:

• If tensor name ends in .weight AND there exist .scales + .biases peers in the shard’s index → this is a quantized weight.
• Load the .weight as MLXArray of dtype uint32 (don’t upcast).
• Load .scales and .biases as their declared dtypes (typically bf16).

3. Extend makeMLXArray for U32

Add a U32 case that builds MLXArray(uint32-buffer, shape) without conversion.

4. Wrap Linears as QuantizedLinear at model construction

The hardest part. Today TinyGPTModelHF creates standard Linear instances. Need:

// At construction time, if quantizedSpec present:
class HFAttention {
    @ModuleInfo(key: "q_proj") var qProj: Linear  // → QuantizedLinear
    ...
}

Two paths:

• A: Always declare as protocol/base Linear, swap concrete instance at load. mlx-swift’s QuantizedLinear IS a Linear subclass, so this works as a value type, but @ModuleInfo wants a concrete type for reflection.
• B: Parameterize the model class on quantization spec, conditionally use QuantizedLinear at construction. Cleaner but requires a second model variant or generics.

Recommended: A, with a post-load step that replaces each targeted Linear with a QuantizedLinear via model.update(modules:). mlx-swift’s Module.update(modules:) allows substituting child modules.

5. Wire scales + biases into the QuantizedLinear’s params

After substituting Linear → QuantizedLinear, write the loaded uint32 weight + bf16 scales + bf16 biases into the new module’s params via model.update(parameters:).

6. Verify forward pass

mlx-swift’s QuantizedLinear.callAsFunction does quantized_matmul automatically. No changes needed in attention/MLP code — they call qProj(x) which dispatches polymorphically.

7. Smoke + accuracy test

• Load v9-LoRA-baked → quantize via mlx_lm convert -q → load via posttrainllm serve.
• Run fm-fixtures-v2 + compose suite.
• Compare to fp16 baseline: accuracy delta should be ≤2% absolute (vs current 33% + 70%).
• Measure: disk size, RSS during inference, TTFW, tok/s.
• Compute formula = (speed × accuracy) / cost; expect ~30% improvement over fp16.

Scope estimate

~1 full day of focused Swift:

• 1 hr: dtype handling + quantization-detection in loader
• 2 hr: model-substitution pattern with mlx-swift Module.update
• 1 hr: scales/biases parameter wiring
• 2 hr: tests + debugging (mlx-swift quirks)
• 1 hr: end-to-end smoke + formula comparison

Compatibility

• v8/v9 fp16 baked-hf: untouched, load same way (no quantization block in config)
• New Q4/Q8 baked-hf: detected via config.json, routed through quantized path
• LoRA on quantized base: separate concern — DoRA/LoRA wrapping over QuantizedLinear. Defer to v2 of this work; v1 supports quantized base only.

Alternative — dequant-at-load (rejected)

Could use dequant_mlx_generic.py to dequant Q4 → fp16 before loading. Saves disk only, no RAM/speed wins. Defeats the formula payoff. Reject.

Alternative — switch Pace to mlx_lm.server (rejected)

mlx_lm has its own serving with native quantized model support. But Pace’s plan per posttrainllm/docs/prds/scope-narrowing-posttrainllm-as-factory.md is embedded MLX-Swift inference (no HTTP at runtime). mlx_lm is Python — doesn’t fit. Reject.

Done when

• Tinygpt serve loads mlx_lm convert -q --q-bits 4 output without error
• Forward pass produces sensible output (smoke test: “open my email” → valid JSON)
• fm-fixtures-v2 score within 2pp of fp16 baseline
• Formula score for Q4 > formula score for fp16 by ≥20%

Related

• feedback-posttrainllm-north-star — the formula
• feedback-pace-decision-framework — accuracy > speed > footprint
• mlx-swift QuantizedLinear reference: github.com/ml-explore/mlx-swift Sources/MLXNN/QuantizedLinear.swift
• Existing dequant tooling: scripts/ane/dequant_mlx_generic.py (the inverse — can re-use for testing)