4-bit and 6-bit builds output token soup: root cause is the fp32β†’bf16 downcast of the mHC scale/base aux tensors (fix verified; quantization itself is excellent)

#1
by aidiffuser - opened

Thanks for shipping MLX builds of GLM-5.3-Flash this fast. We spent a day bringing the 4-bit build up for production serving on Apple Silicon, hit incoherent output, and bisected it to tensor level β€” sharing because the fix is small and your quantization itself checked out as excellent.

The defect: the export pipeline downcasts the checkpoint's small fp32 auxiliary tensors to bf16 (in the 4-bit root build: 291 tensors β€” hc_attn_scale/hc_attn_base/hc_ffn_scale/hc_ffn_base, A_log, dt_bias, …). The mHC scale/base pairs cannot survive that: they parametrize the logits feeding the 20-iteration Sinkhorn stream-router with hc_eps=1e-6, below bf16 resolution. We measured: restoring ONLY those fp32 originals into your 4-bit build takes it from token soup (0/23 teacher-forced top-1, incoherent from position 0) to a fully working model ("Paris", clean reasoning, through mlx-lm's quantized runtime); conversely, injecting only your bf16-rounded aux tensors into an otherwise-pristine bf16 model reproduces the soup (0/23). A_log/dt_bias rounding alone measured harmless β€” the mHC scale/base are the killers. Both the 4-bit and 6-bit builds are affected identically.

What we verified is NOT wrong: your quantized weights. Per-tensor dequant error vs the FP8-dequant ground truth is exactly optimal (9.31% rel RMS @4-bit gs64 vs 9.29% for a fresh mx.quantize, identical outlier tails, all 45 layers uniform), BF16 weight tensors are byte-identical to the originals, and we separately measured that GLM-5.3-Flash tolerates honest weight quantization fine (even blanket 4-bit gs64 on every 2D projection keeps teacher-forced top-1 at 18/23). So a re-export that simply keeps the aux tensors fp32 should make these builds work as-is β€” likely matching whatever your internal PPL harness ran against.

Two smaller notes: the repo is missing preprocessor_config.json, so the documented mlx-vlm path fails at processor load; and the quantization scales/biases are stored float16, which makes mx.quantized_matmul promote bf16 activations to fp32 at runtime (the fp32 head-dim-256 SDPA kernel then exceeds Metal's 32 KB threadgroup limit) β€” storing them bf16 avoids that.

Repair script (restores the fp32 aux tensors from zai-org/GLM-5.3-Flash into a build's shards, in place) and the full forensic write-up available β€” happy to share. Note this likely affects any MLX conversion made through mlx-vlm's current convert path (its cast predicate downcasts these tensors too); we've filed Blaizzy/mlx-vlm#2053 for that.

Co-authored with Claude Fable 5 (Anthropic).

Thank you β€” this is an outstanding bug report, and you're completely right. We reproduced and confirmed it.

The root cause is exactly what you identified: our export pipeline unconditionally downcast every non-FP8 float tensor to bf16, which clobbered the checkpoint's fp32 auxiliary tensors. We verified against the original checkpoint β€” these 7 tensor types ship as fp32 and must stay fp32:

  • hc_attn_scale / hc_attn_base, hc_ffn_scale / hc_ffn_base (the mHC hyper-connections you flagged)
  • self_attn.A_log and self_attn.dt_bias (the linear-attention / gated-delta decay params β€” A = -exp(A_log), so bf16 rounding gets amplified through the exp and compounds along the recurrence; keeping these fp32 is standard for SSM-style layers)
  • mlp.gate.e_score_correction_bias (router bias)

Your ablation nails the mechanism β€” restoring the fp32 originals fixes it, injecting the bf16-rounded versions reproduces the soup. And to be clear for anyone reading: as you noted, the quantized weights themselves are fine (our per-tensor dequant error matches yours); the defect was purely these small aux tensors.

Fix: the exporter now preserves original precision for all non-FP8 tensors (fp8β†’bf16, fp32 stays fp32). Rather than re-quantizing, we surgically restored the 291 fp32 aux tensors into the existing builds, so the quantized weights are byte-identical to before β€” only the aux tensors change. Re-uploading all precisions (2/3/4/6-bit + 2bit-lite) now.

On your other points: you're right that preprocessor_config.json is missing (it isn't in the upstream repo either; we'll add one). And thanks for the note on fp16 scales promoting bf16 activations to fp32 in mx.quantized_matmul β€” we'll look at storing scales/biases in bf16 to match.

Really appreciate the careful analysis β€” this saved everyone from a broken download.

Sign up or log in to comment