fix(qwen3): quantize Model before loading safetensors for 4-bit weights #2
Loading…
Reference in a new issue
No description provided.
Delete branch "fix/qwen3-quantize-before-load"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
Cherry-picks an already-authored, already-verified fix (originally commit
3b607a72, 2026-05-14) that never made it tomain— the branch it lived on diverged and was never merged.Bug
Qwen3 checkpoints that declare
quantizationinconfig.json(themlx-community/Qwen3-*-{4,8}bitfamily) buildModel::newwith every projection/embedding asMaybeQuantized::Original.load_qwen3_modelthen loads packed[..., dim/8]safetensors weights straight into thoseOriginalslots, which expect[..., dim]. The result is a tensor with the wrong last dimension, and the firstRMSNormin the forward pass crashes:Fix
Before
load_safetensors, ifconfig.jsondeclaresquantization, callModel::try_into_quantized(group_size, bits)first — converting everyMaybeQuantizedfield to itsQuantizedvariant so the packed weights land in slots that actually invokemlx_dequantize.Verified
mlx-community/Qwen3-0.6B-4bit: previously crashed on the first generate step with the exact error above; now generates real tokens.ritual_engineintegration tests (mlx_concat_smoke,mlx_smoke) against a real local model snapshot — both now pass end-to-end (tokenizer → generate loop → decode), where they previously failed identically with the same rms_norm crash whether usingConcatKeyValueCacheor the customRitualLayerCache, confirming the bug was in this model path, not in downstream KV-cache wrapping.Rebased cleanly onto current
main(cherry-pick, no conflicts).mainalready independently picked up the.inner-qualified key fallback fix from the same branch (via a different commit), so this PR only adds the Qwen3-specific piece that's still missing.When a Qwen3 checkpoint declares quantization in config.json (the mlx-community/Qwen3-*-{4,8}bit family), Model::new builds the model with every projection / embedding wrapped in MaybeQuantized::Original. load_qwen3_model then calls load_safetensors against those Original slots — silently loading packed [vocab_size, hidden_size/8] tensors into Embedding slots whose Module::forward returns a tensor with the wrong last dimension. The first RMSNorm in the forward path then errors with: [rms_norm] (*weight) must have the same size as the last dimension of x but has 2048 elements (or 1024 for Qwen3-0.6B) because the actual x carries the un-unpacked packed dimension rather than hidden_size. Fix: - Add `quantization: Option<QuantizationConfig>` to ModelArgs, populated via serde from config.json's `quantization` block (Optional so bf16 checkpoints continue to deserialize unchanged). - Before load_safetensors, if `quantization` is present, call `Model::try_into_quantized(group_size, bits)` on the freshly-built model. That converts every MaybeQuantized<nn::Linear / nn::Embedding> field to MaybeQuantized::Quantized so the packed weights now land in slots whose forward path actually invokes mlx_dequantize. Verified end-to-end: - mlx-community/Qwen3-0.6B-bf16: still loads + generates (no regression on the unquantized path). - mlx-community/Qwen3-1.7B-4bit: previously crashed on layer 0 input_layernorm; now reaches the generate loop and produces tokens. Test reproducer (downstream): see ritual_engine `mlx_concat_smoke` integration test pinned at this rev — was 1 failed, now 1 passed.