fix(module): fall back to .inner-qualified key when loading quantized safetensors #1

Merged

peterswimm merged 1 commit from fix/load-safetensors-inner-key-remap-onmain into main

2026-07-12 23:12:45 +00:00

forge-ci commented

2026-07-12 21:50:09 +00:00

Member

Summary

QuantizedEmbedding/QuantizedLinear wrap their real weight in a #[param] inner: Embedding/inner: Linear field, so the flattened Rust param path is <path>.inner.weight, while mlx-community checkpoints store it flat as <path>.weight.
Module::load_safetensors only did an exact-match lookup with no else/error branch, so any such key silently matched nothing -- the quantized weight stayed at try_into_quantized's throwaway random-init value (correct shape/dtype, so it loads and runs without error), while scales/biases (whose names have no such prefix mismatch) loaded correctly. Net effect: every quantized linear/embedding in a loaded model produced plausible-shaped but semantically wrong output.
Root-caused this against the Qwen3-Embedding-8B-4bit-DWQ checkpoint by comparing raw dequantize() output (correct, bit-exact vs. a known-good Python mlx_embeddings reference) against the full model's output (wrong) -- isolating the bug to weight loading, not the dequantization math.
Fix: when the exact key has no match, retry with .inner inserted before the final segment before giving up.

Test plan

cargo test -p mlx-rs -- new inner_qualified_key unit tests pass
Downstream verification: a Qwen3-Embedding-8B-4bit-DWQ forward pass through mlx-lm's load_qwen3_model, before vs. after this fix -- cosine similarity between semantically related sentences went from ~0.02 (near-random) to 0.73 (matching a Python mlx_embeddings reference to 3 decimal places)

## Summary - `QuantizedEmbedding`/`QuantizedLinear` wrap their real weight in a `#[param] inner: Embedding`/`inner: Linear` field, so the flattened Rust param path is `<path>.inner.weight`, while mlx-community checkpoints store it flat as `<path>.weight`. - `Module::load_safetensors` only did an exact-match lookup with no else/error branch, so any such key silently matched nothing -- the quantized weight stayed at `try_into_quantized`'s throwaway random-init value (correct shape/dtype, so it loads and runs without error), while `scales`/`biases` (whose names have no such prefix mismatch) loaded correctly. Net effect: every quantized linear/embedding in a loaded model produced plausible-shaped but semantically wrong output. - Root-caused this against the `Qwen3-Embedding-8B-4bit-DWQ` checkpoint by comparing raw `dequantize()` output (correct, bit-exact vs. a known-good Python `mlx_embeddings` reference) against the full model's output (wrong) -- isolating the bug to weight *loading*, not the dequantization math. - Fix: when the exact key has no match, retry with `.inner` inserted before the final segment before giving up. ## Test plan - [x] `cargo test -p mlx-rs` -- new `inner_qualified_key` unit tests pass - [x] Downstream verification: a Qwen3-Embedding-8B-4bit-DWQ forward pass through `mlx-lm`'s `load_qwen3_model`, before vs. after this fix -- cosine similarity between semantically related sentences went from ~0.02 (near-random) to 0.73 (matching a Python `mlx_embeddings` reference to 3 decimal places)

forge-ci added 1 commit

2026-07-12 21:50:09 +00:00

fix(module): fall back to .inner-qualified key when loading quantized safetensors

validate / tests (macos-15, 1.85.0) (pull_request) Failing after 4s

Details

validate / tests (macos-15, stable) (pull_request) Failing after 2s

Details

validate / checks (pull_request) Failing after 3s

Details

46da0d0c4f

QuantizedEmbedding/QuantizedLinear wrap their real weight in a #[param]
inner: Embedding/Linear field, so the flattened Rust param path is
<path>.inner.weight, but mlx-community checkpoints store it flat as
<path>.weight. load_safetensors only ever did an exact-match lookup with
no else/error branch, so any such key silently matched nothing -- the
quantized weight stayed at try_into_quantized's throwaway random-init
value (correct shape/dtype, so it loads and runs without error), while
scales/biases (whose names have no such prefix mismatch) loaded correctly.
Net effect: every quantized linear/embedding in a loaded model produced
plausible-shaped but semantically wrong output.

Root-caused by comparing raw dequantize() output (correct, bit-exact
against a known-good Python mlx_embeddings reference) against the full
model's output (wrong) for the Qwen3-Embedding-8B-4bit-DWQ checkpoint --
isolating the bug to weight loading, not the dequantization math.

Fix: when the exact key has no match, retry with .inner inserted before
the final segment before giving up.

peterswimm merged commit dcda6bef1b into main

2026-07-12 23:12:45 +00:00

peterswimm deleted branch fix/load-safetensors-inner-key-remap-onmain

2026-07-12 23:12:46 +00:00

peterswimm referenced this pull request from a commit

2026-07-12 23:12:47 +00:00

Merge pull request 'fix(module): fall back to .inner-qualified key when loading quantized safetensors' (#1) from fix/load-safetensors-inner-key-remap-onmain into main