feat(qwen36): load MLX-affine dense weights and norms from a qpack container - #1343
Open
Avicennasis wants to merge 10 commits into
Open
Avicennasis wants to merge 10 commits into
Avicennasis wants to merge 10 commits into
Conversation
load_tokenizer indexed the BPE merge table from tokenizer.json by splitting each entry on a space: the legacy "a b" string spelling. tokenizers >= 0.20 (transformers 4.45+) writes each merge as a two-element array ["a","b"] instead, and that is what the Qwen3.6 checkpoints ship. Every entry failed the string test, `continue` skipped it, and the reader finished with ZERO merges and no message. Nothing downstream treats an empty merge table as an error: bpe_piece reads it as "nothing to merge", so encode_text degraded to one token per byte-symbol. Measured on a real-model run before the fix: a 24-character prompt became 24 tokens, and the model was fed a sequence it had never seen in training. Accept both spellings. A string entry splits on the first space as before; a two-element array of strings takes its halves directly; any other shape is skipped as it was. Both forms index into the same smap_put(&g_merge, key, r) table with the same 0x1F-joined key, so the rest of the tokenizer is unchanged. tests/test_qwen36_tok_merges writes two tokenizer.json files for the same five-piece vocabulary that differ only in the merge spelling and asserts that both encode "in in" to the same merged ids [in, Ġin] -- one plain merge (i+n) and one chained merge (Ġ+in) whose input the first produces. On the pre-fix reader the pair form loads zero merges and the encoding comes back as one id per byte-symbol, so the gate fails. The scratch directories the test binaries write (tests/tmp_*) are added to .gitignore. (cherry picked from commit d783505)
The Darwin arm of the Makefile finds libomp with `brew --prefix libomp`. On a CLT-only Mac driven non-interactively (`make` from a script, over SSH, in CI) Homebrew is routinely installed but not on PATH, the probe prints nothing, and the OpenMP arm never engages. Nothing warns: the build is a valid single-threaded engine with libomp sitting right there on disk. Measured on an M4 Max with the Qwen3.6-35B qpack: 0.87 tok/s from the silent single-threaded build, 4.44 tok/s (range 3.80-4.76, thermal) once OpenMP engaged, TTFT 3.89 s to about 0.7 s, output byte-identical by cmp. Engine code is unchanged by this commit -- every omp site is row-partitioned and deterministic by construction -- so the difference is entirely the missing flag. When the brew probe comes back empty, fall back to the two standard Homebrew prefixes (/opt/homebrew/opt/libomp on Apple silicon, /usr/local/opt/libomp on Intel). The existing artifact check still runs after the fallback: both omp.h and libomp.* must exist before any flag is added, so a bare prefix directory without the library still yields the plain build, as before. (cherry picked from commit 44ddb39)
…ract affine_quant.h carries one reading of MLX's packed affine words: coli_affine_matmul_ref folds scale*q + bias into its accumulation and never materialises the weights. That is the right shape for a routed expert dispatched per token, and the wrong shape for the dense tensors the qwen36 engine keeps RESIDENT -- attention, DeltaNet, shared expert, norms, embeddings -- which a Swiftlet qpack container also stores as affine triples in its model.safetensors and which the engine wants as f32 rows exactly once, at load. coli_affine_dequant_ref expands a validated view into row-major out[output_dim, input_dim] = scale[row, col/gs] * q[row, col] + bias[row, col/gs], with the same bit conventions as matmul_ref (little-endian uint32 words, lowest bits = lowest logical column, Q4 or Q8, BF16/F16/F32 scalars). It runs coli_affine_validate first, so a malformed view comes back as a status and never as plausible output. It lives next to matmul_ref deliberately: the two bit-level readings of the same words cannot drift apart without this file changing. tests/test_affine_dequant pins the contract from the bits up: a Q4 word authored as 0x76543210 must expand to logical columns 0..7 in that order (a most-significant-first reading produces the reversed sequence and fails the exact-value checks); the affine map is checked with values exactly representable in BF16, F16 and F32 so the comparisons are equalities, not tolerances; a crafted Q8 word proves the byte order; on a multi-group view x @ dequant(W)^T must equal matmul_ref(x, W); and NULL buffers, truncated scales, and a group size that does not divide input_dim are refused with statuses.
A Swiftlet qpack container carries the dense half of the model in its model.safetensors exactly as mlx-lm wrote it: every quantized tensor is a packed U32 `.weight` with `.scales`/`.biases` siblings (BF16 on the production container), and the text stack sits under a `language_model.` prefix because the checkpoint is multimodal. qwen36's loader could do nothing with that file -- st_init exit(1)'d on the first U32 dtype string before the rest of the header was even indexed, and load_t_n only speaks unprefixed names. Running the engine over a container therefore meant re-converting weights whose bytes the container already holds. Three additive changes let the loader read the container as it is: st.h indexes U32/I32 as dtype 7 (4 bytes/element), the same way the fp8 codes were added: the header parses to the end, the tensor is findable with its packed shape, and the float readers keep refusing it by value (dtype >= 3 in st_read_f32) -- callers must go through st_read_raw with geometry they validated. dense_resolve tries the plain name first and `language_model.<name>` second, so converted snapshots are untouched and the container's spelling is found; when neither exists it returns the CANONICAL name so the refusal names the tensor the engine wanted. load_t_affine expands a U32 triple through coli_affine_dequant_ref. Every dimension is anchored to `want`, the config-implied element count the forward pass will index with -- the discipline load_t_n already applies to plain tensors -- so a container whose packed geometry disagrees with config.json is a refusal, never a plausible heap OOB. Bits (Q4 = 8 per word, Q8 = 4) and the group size are DERIVED from the shapes (logical input over packed words, logical input over scale groups) rather than parsed from config quantization overrides: the file's own byte layout is what the expansion must agree with, and coli_affine_validate re-checks the derived geometry against every buffer length. Scales in BF16/F16/F32 are accepted; any other dtype, mismatched biases, an orphan `.weight` without siblings, a `want` the packed shape cannot tile, or a per-word count that is neither 4 nor 8 exits with a message naming the tensor. load_t_n dispatches to load_t_affine on dtype 7 and is otherwise unchanged; the optional-tensor probes (q/k norms, router bias, shared expert gate) go through dense_has so the prefixed spelling is found too. The router correction bias now loads through load_t_n as well, which gives it the size check the other dense tensors already had. tests/test_qwen36_dense_affine builds a safetensors file by hand (no numpy, no torch): both name spellings load and the plain one wins when both exist; st_init indexes the U32 tensors with the right rank and shape while st_read_f32 still refuses them; a Q4 and a Q8 triple with authored bit patterns expand to exactly scale*q + bias; and (POSIX fork gates) an untileable want, an orphan U32 weight, and a wrong-size plain tensor each exit non-zero instead of returning rows. With argv the same binary loads one named tensor through load_t_n and compares it against an expected.f32 file, which the snap-view tool test uses as its cross-implementation round trip.
…d form rmsnorm_row applies its weight as (1 + w). That is the HF Qwen3.6 convention -- the norm weights are stored zero-centered -- and every snapshot convert_qwen36.py produces keeps it, so the forward pass was written for it and has only ever seen it. mlx-lm materialises the +1 into the weights at conversion, so an MLX-derived container stores FULL gamma. Measured on the production Qwen3.6-35B qpack: input_layernorm mean 1.03, q_norm mean 1.33, where the zero-centered forms centre on 0. Fed through (1 + w) that doubles every normalised activation and the model degenerates to noise, with every tensor loaded and every shape check passing. The shift is undone at LOAD, so the forward pass keeps exactly one convention. Cfg gains zero_centered_norms (default 1, the HF dialect; qwen36_meta.json may set it false, which is what the snap-view tool emits for a container). load_norm_n loads through load_t_n and, when the flag is 0, subtracts 1 from every element. It is used for precisely the tensors rmsnorm_row touches: input/post layernorms, q/k norms, and the final norm in both the main loader and the edge engine. The MoE router weight, which shared the LD macro with the layernorms, is loaded separately and never shifted. The DeltaNet gated norm is full gamma in BOTH dialects (its forward multiplies plain w) and does not go through this loader. tests/test_qwen36_dense_affine gains two gates: load_norm_n leaves the values untouched under the default dialect and returns exactly value-1 under the MLX one, and load_meta parses a zero_centered_norms:false qwen36_meta.json into the config.
The engine's loading contract is the one convert_qwen36.py established:
a snapshot directory with a FLAT config.json, a qwen36_meta.json whose
head dims were derived from the actual weight shapes, the *.safetensors
holding the dense weights, and tokenizer.json. A Swiftlet qpack
container already carries every byte the engine needs --
model.safetensors (the dense half, which the loader now expands in
memory), tokenizer.json, and a nested multimodal config.json -- but not
the two engine-side metadata files, and the container itself must never
be written: it is a production artifact shared with Swiftlet and
hash-pinned by hashes.json.
So the bridge is a VIEW. tools/make_qwen36_qpack_snap.py writes a fresh
directory holding a generated flat config.json and qwen36_meta.json next
to symlinks into the container. No tensor bytes are copied or converted
on disk; st_init follows the link and the affine dense loader does the
expansion at load. Stdlib only: the safetensors header is 8 bytes of
length plus JSON, so there is no torch, numpy or safetensors dependency.
An --out inside the container is refused.
Head-dim derivation mirrors convert_qwen36.py ("derived from the actual
weight shapes, authoritative") with one extra step: quantized tensors
declare PACKED shapes (uint32 words along the input axis), so logical
dims are unpacked through the config's MLX quantization spec -- default
bits/group plus the per-module overrides, matching what Swiftlet's
Checkpoint.quantSpec applies. The meta declares zero_centered_norms:
false, because an mlx-lm container stores full-gamma norm weights and
the engine unshifts them at load on that flag. The argparse description
falls back to a literal under python -OO, where __doc__ is None.
python tools/make_qwen36_qpack_snap.py \
--container ~/models/qwen3.6-35b.qpack --out ~/build/qwen36-35b-snap
SNAP=<out> QWEN36_QPACK=<container> ./qwen36 16 4 prompt.txt
tests/test_make_qwen36_qpack_snap.py builds a synthetic container
(nested config with per-module overrides, a hand-written safetensors
with packed U32 triples) and checks the flattened config, the unpacked
head dims (a derivation that forgets to unpack reports o_in 8x too small
and fails), layer_types / n_active / DeltaNet dims / rope_theta
surviving the trip, the links, the untouched container, the refusal, and
the dialect flag. The round trip is the cross-implementation evidence:
the fixture carries a seeded-random Q4 tensor whose dequantized rows are
computed in pure Python (bf16 decoded exactly, scale*q + bias per group)
and handed to tests/test_qwen36_dense_affine's compare mode, which loads
the same bytes through the engine's own st_init + load_t_n path via the
view's symlink. It skips loudly if that gate binary is not built.
Contributor
Author
|
Apple/Metal gates at |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this adds
The loader half of what was #1323:
qwen36reads a Swiftlet-format qpack container's dense weights and norms from the container's ownmodel.safetensors, through #1290's affine contract. Nothing here touches Metal; the routed experts are unchanged, and the Metal slot pool that consumes them follows in #1323, which is now stacked on this PR.Four commits of our own:
feat(affine): a loader-side dequant oracle in the affine contract (coli_affine_dequant_ref, same bit conventions as the matmul reference), withtests/test_affine_dequant.c.feat(qwen36): dense weights from the container.st.hgains U32 as an additive dtype;load_t_affineexpands the MLX-affine U32 triples at load with geometry derived from the shapes and re-validated against the config (a mismatch is a refusal, never a plausible out-of-bounds read).tests/test_qwen36_dense_affine.c.feat(qwen36): norm dialect declared, not guessed. mlx-lm materialises the+1in the RMSNorm gammas the engine applies itself; azero_centered_normsmeta flag (default = the HF convention, so converted snapshots are unchanged) unshifts once at load.tools(qwen36):make_qwen36_qpack_snap.py(stdlib-only) builds the flatconfig.json+ shape-derivedqwen36_meta.jsonsnap view with symlinks into the read-only container, following theconvert_qwen36.pycontract ("engine reads flat config + meta; tooling stays smart");tests/test_make_qwen36_qpack_snap.py.Own diff: 8 files, +1221/-16 on top of the prerequisites.
Stacked on
Their commits show in this diff until they land (16 of the 21 files, +2437/-35). The branch is a replay of their current heads (
7f3fe6c9,cfd51066,f01621ff) ontodevatf58a2679; each replayed commit has the same patch-id as its source (git range-diffflags #1319/#1320 on hunk context only).Validation
Linux (x86-64, gcc), this branch at
05868e5b:make -C c check(clean, portable build,test-c+test-python): exit 0.test-cran 112 gate binaries, none failed;test-python739 tests OK, 35 skipped.make -C c test-asan: exit 0, the same 112 binaries clean under ASan + UBSan.Nothing in these four commits is Metal-specific, so CI's macOS
make checklane compiles and runs the same gates as Linux. Not re-run in this pass (no Apple machine): the round-trip of the production container's dense tensors against an independent pure-Python dequant, reported asmax |diff| 0in the original #1323 description on the pre-rebase series.Run recipe
builds the snap view the engine reads through
SNAP=<dir>. Running a qpack container end to end also needs the routed-expert path (QWEN36_QPACK=<dir>,METAL=1), which is #1323.