Repository navigation
llama: published MLX affine quants (8-bit and 4-bit) read in place, beside the bf16 default (builds on #1415) - #1416
BobClawblaw wants to merge 7 commits into
Conversation
|
Nothing converts at load, which is right. But when Also: |
73e9b7f to
0500646
Compare
|
Thanks, agreed. Pushed 0500646 (one commit on #1415's head f92916d, which is rebased onto main 15a0b2a): MLX's own layout, read in place. Checked on two published community quants, pinned:
The engine is as close to the release as MLX is on each quant. 16- and 7-row windows and one-row steps agree bit for bit on both. Every harness serving check passes on both: structured 18/18, tools 5/5, sampling 9/9, turns, soak 60/60, concurrent 12/12, streaming, multi-turn, stop, thinking, max_tokens. Contract. Capabilities. bf16 A/B. The The PR body is rewritten for the new scope (published quants, no conversion). On these commits over main a1e6d34: |
BobClawblaw
changed the title
|
Thanks @BobClawblaw. The published split MLX layout is now read in place ( Two asks:
#1415's second point (the exactness mode always exits 0) applies to this branch too. |
A new `llama` family for the native Metal engine: dense Llama-architecture checkpoints (heads of 128, default RoPE, no biases, an untied head) in bf16, read in place from the shards with nothing converted. It keeps HF modeling_llama's rounding points (a bf16 residual, LlamaRMSNorm, RoPE with bf16 cos and sin, SiLU), runs prompt chunks on gemm_nax's tensor-unit GEMM and decode windows on a simdgroup-a-row matvec, decode attention on sdpa_vector lanes at absolute positions, and sampled draws on tf_sample_full. tf-llama-run checks greedy tokens, every position's logits, and that windows and one-row steps agree bit for bit. Measured on an M5 Ultra against HF's bf16 run of the release: 1,500 positions teacher-forced at mean KL 0.0015 and top-1 99.1%; decode 16.4 tokens/s a stream and 28-32 with two; cold prompts 1.3k tokens/s at 8k tokens and 563 at 131k. docs/recipes/nanbeige-4.3-33b-metal.md has the rest.
The family's Metal kernels through gates 1-5 (zig/tests/contract/metal_llama*.zig, the roll call and sources in contracts.zig): the glue (embedding, the norms at HF's rounding points, the residual add, RoPE with bf16 cos and sin, the SiLU gate, the decode attention's partials and their merge, the argmax draw by its stated NaN rule), the six row-projection forms, the four GEMM forms and the prompt attention at both served geometries, each with its float64 reference, served shapes and tiles, launched as forward.zig launches it; gate 5 reads the pipelines on this GPU. tf-kernel-contract --only llama on the M5 Ultra: 21 kernels, 21 pass gates 1-4, 0 fail, 0 known findings, 0 limit problems.
…out tensor units by name; the teacher-forced misses and the licence in the doc tf-llama-run takes exact's window width as an argument and loses its environment switches (LLAMA_LAYERS, LLAMA_WINDOW).
…hecks a planted row is reported
… near-tie, a float64 reference, the doc's misses and a one-process receipt
…eside the bf16 default
The family reads an MLX affine checkpoint as mlx_lm writes it: each projection
its own {u32 words, bf16 scales, bf16 biases} triple in groups of 64 at 4 or 8
bits, the embedding and head quantized, the norms bf16, and config.json's
uniform quantization table. Nothing is converted, by us or at load. Decode
windows run the core's row kernels, prompt chunks its tensor-unit affine
matmul; the embedding's rows come from a new family kernel, ll_embed_affine,
MLX's dequantize bf16(bf16(scale * q) + bias), registered with the contract.
affine_mm's 8-bit groups-of-64 form is registered too. capabilities lists
bf16, mlx-q4g64 and mlx-q8g64. The sums scratch is allocated only for
quantized checkpoints, so the bf16 path is as before.
Nanbeige4.3-33B on an M5 Ultra, two published quants pinned
(pratikgx/Nanbeige4.3-33B-MLX-4bit, fcmeyer/Nanbeige4.3-33B-MLX-8bit),
against HF's bf16 run of the release over 1,500 positions: 8-bit mean KL
0.00234, top-1 98.7% (MLX on the same files: 0.00272); 4-bit 0.0849, 91.1%
(MLX: 0.0851). Windows and one-row steps agree bit for bit on both; the
release's logits are bit-identical to before. The recipe doc has the tables;
the quants are lossy and the release stays the default.
…wo paths and a float64 pass, and a one-process receipt
0500646 to
138b785
Compare
|
Thanks. Pushed 138b785 (one commit on #1415's new head): 1. Every differing choice on both quants.
The recipe now has a table per quant with every position where the engine differs from any of them: each reference's logits for the contenders and both verdicts. Causes. At every red position the exact margin between the contenders is under 0.46: 8-bit 401 0.39, 425 three tokens within 0.13, 1370 three within 0.16, 1447 0.45; 4-bit 533 0.33, 600 0.42, 1355 three within 0.45, 1418 three within 0.26. Every bf16 path carries more error than that. Against the float64 pass, at each position's exact top 8 tokens and less the row's common shift, the root-mean-square logit error over all 1,500 positions is:
MLX's own paths are red against the float64 pass at the same kind of positions: on the 8-bit quant at 425 and 1461 (one pass) and 1370 and 1447 (decode), on the 4-bit at four (one pass) and five (decode), four of them shared with the engine's. The largest engine-MLX difference, 8-bit 1461 ( The greedy chats: they part from MLX's on four of six chats (8-bit) and three (4-bit). At each parting token, both sides' logits as each greedy run reaches them (the engine's chunks and one-row steps;
At both red ones the float64 pass picks the engine's token (4-bit chat1: 28746 37.7532 against 570 37.4911; 8-bit chat2: 152404 26.5785 against 145762 26.4806). MLX rounds each pair to a tie, and its argmax takes the lower id. Row 1447 (red on the 8-bit quant against float64, where MLX's decode ties the same two tokens) has a larger engine error than most, so I traced it on the bf16 release, where it is a near-tie. I stopped the engine's forward after each layer and after each layer's attention (a local probe, not in the PR), and compared with HF's hidden states and the float64 pass. The engine's residual tracks HF's error through layer 35. At layers 36-40 its attention sublayers add 1.7-4 times HF's error and its MLPs grow theirs by 18-21% against HF's 7-10%. After that both grow at the same rate, ending at 12.4% (engine) against 4.8% (HF) relative error at layer 51. The excess sits mostly in the massive-activation channel 4086 (750-1,700 over those layers, a bf16 step of 4-8). Rounded at the head, the engine ties
2. A clean receipt, no other model server or GPU job on the Mac, this branch's build (a CPU float64 reference run overlapped its first minute):
Token idsCI on these commits over main 03d08f1 (the pushed heads are rebased onto d0c9dc9, which adds only a CI workflow time limit): on #1417's head, which carries all three, |
|
Thanks @BobClawblaw. On What's left is the teacher-forced failures. Seven positions still fail against float64: 425, 1370 and 1447 at 8 bits, and 533, 600, 1355 and 1418 at 4 bits. The trace through layers 36-40 and channel 4086 stops before it names the operation that adds the error, and an average over all 1,500 positions can't clear a failing row. Please find that operation and correct it, keeping the current contender and margin rule, and give each affected row's own error. Until those rows pass, the quants can't land. The non-finite request for the checker on #1415 applies here too. |
|
Thanks. We looked for the operation, and there isn't one to correct. Rather than ask for an exception to the rule, we're withdrawing the quants: closing this PR, and #1417 now builds on #1415 alone, with its drafter served from the bf16 release. Here's what we found, for the record. Per operation. A local probe (not in any PR) stopped the engine's forward at row 1447 at layers 36, 38 and 40, on the bf16 release and on the 8-bit quant, and kept every operation's output. A float64 checker then recomputed each operation from the engine's own input buffers, with the reference's rounding points and the quant's exact weights (scale * q + bias). For each one: q|k|v, attention, o_proj, both residual adds, gate, up, silu(gate) * up and down, the engine's output is within half a bf16 step of the exact result at the vector's scale (0.06 to 0.50 steps), on both weight sets. Each operation is correctly rounded given its inputs. Elementwise, a few values sit one step off where the fp32 sum and the exact sum fall on opposite sides of a rounding boundary. A perfect implementation of the rounding points. The other three (8-bit 1370 and 1447, 4-bit 600) pass for the perfect implementation and fail for the engine. The difference is the fp32 accumulation inside the projections, the one-step flips above, compounded over 52 layers. MLX's GPU decode fails the same rule at 1370, 1447 and 600 for the same reason. Matching the perfect implementation bit for bit would take exact sums in every projection. Each failing position's own error (float64 logits of the three contenders; the engine's; the largest gap):
If you later want lossy community quants judged against a reference with the same rounding points, the code is on our fork's |
What
The
llamafamily serves published MLX affine quants (8-bit and 4-bit, groups of 64) of a Llama release in place, asmlx_lmwrites them, beside its bf16 default (builds on #1415). Nothing is converted, by us or at load. Checked on two community quants of Nanbeige4.3-33B, pinned. They are lossy; their cost against the release is measured below, and the bf16 release stays the recipe's default.How
weights.zig:Linear = union{bf16, affine}. Each of q, k, v, o, gate, up and down is read as its own{weight (u32 words), scales, biases}triple, and the embedding and head asmlx_lmquantizes them. Shapes, dtypes and byte extents are checked at load.config.zigadmits a uniformquantizationtable (affine, groups of 64, 4 or 8 bits) and refuses anything else by name.row_projection.zig), and prompt chunks the core's tensor-unit affine matmul (affine_mm.zig), as the other affine families do.ll_embed_affine(zig/kernels/metal/llama/glue.metal), which computes MLX's dequantizebf16(bf16(scale * q) + bias).ll_embed_affinehas gates 1-4 at 4 and 8 bits (zig/tests/contract/metal_llama.zig) and a limits entry.affine_mmat 8 bits in groups of 64, a shape the published 8-bit quants use, is registered inmetal_main.zigand the registry.tf-kernel-contract --only llama:SUMMARY 23 kernels: 23 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems;affine_mmq8 g64 passes gates 1-4.capabilities --jsonlistsbf16,mlx-q4g64andmlx-q8g64forllama.sumsscratch is allocated only for quantized checkpoints, so the bf16 path is as in llama: Nanbeige4.3-33B on Metal, the bf16 release served as published, every kernel in the contract #1415.Both state Apache-2.0, like the release.
Cost against the release (M5 Ultra; the HF reference and harness plans of #1415)
The reference is HF's bf16 run of Nanbeige/Nanbeige4.3-33B @ c1093c501cc6460427c0ac634db7223d0450c0ee over 1,500 teacher-forced positions. MLX is MLX 0.32.4 with
mlx_lm0.32.2 running the same quant files.The engine is as close to the release as MLX is on each quant. 16- and 7-row windows and one-row steps agree bit for bit on both. Every harness serving check passes on both: structured 18/18 (drafted == plain 12/12), tools 5/5, sampling 9/9, turns resumed 5,493 of 5,649 tokens, soak 60/60, concurrent 12/12, streaming, multi-turn, stop, thinking, max_tokens.
Receipt
--context 131072 --parallel 2 --prompt-cache-gib 8andLLAMA_LOAD_LIMIT_GB=150, with an idle server of another model loaded beside it.tools/bench_concurrent.py --levels 1,2 --tokens 128 --temperatures 0,0.7 --alone --serialon each quant: every reply equals its request alone and with"draft": false(alone 12/12, serial 12/12, failed 0, on each).tools/bench_openai.py --tokens 256 --reps 3, bf16 → 8-bit → 4-bit:bench_concurrent, greedy: 31.3-34.4 → 49.8-50.4 → 76.4-76.8tools/bench_cache.py --require-cold, a 9,000-token prompt cold: 7.18 s (8-bit) and 6.64 s (4-bit). Repeated raw-token prompts inbench_cache.pyreportcached_tokens0: the family's prompt cache resumes chat turns, and the harness's turns check resumes 5,493 of 5,649 tokens in 0.24-0.28 s against 4.15-4.49 s cold. The cold first request is what the receipt asks for.llama_teacher_judge.pyagainst MLX's one pass, MLX's decode steps and a float64 pass on the quant's exact weights: on the 8-bit quant 1,481 / 1,484 / 1,478 of 1,500 equal (red at 315, 1370, 1447, 1461 / 401, 425 / 425, 1370, 1447), on the 4-bit 1,477 / 1,482 / 1,479 (red at 1321 / none / 533, 600, 1355, 1418). At every red position the exact margin is under 0.46; the engine's logit error against the float64 pass (0.126 and 0.141 root-mean-square at the top tokens) is MLX's (0.117-0.133); MLX's own paths are red against it at 2-5 positions each. The recipe has both tables and the greedy chats' partings.exact0 differing rows; on the 4-bit quant, concurrent == solo for six prompts and resumed == fresh for three chats (8,521-8,528 tokens kept).zig build test,test-kernel-contract,test-golden,fuzz-cortheon,dev-tools,test-server-securityandtest-pull-securityall exit 0, and on llama: Nanbeige4.3-33B on Metal, the bf16 release served as published, every kernel in the contract #1415's and llama: published MLX affine quants (8-bit and 4-bit) read in place, beside the bf16 default (builds on #1415) #1416's headszig build testandtest-kernel-contractexit 0;zig fmt --checkis clean and the lean check finds 0 new problems (428 on main and on each head);llama_teacher_judge_test.pypasses;tf-kernel-contract --only llamaon the M5 Ultra:SUMMARY 23 kernels: 23 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems(21 kernels on llama: Nanbeige4.3-33B on Metal, the bf16 release served as published, every kernel in the contract #1415's head). Rebuilt on the rebased heads, the bf16 release's 1,500-position logits are bit-identical to the judged run on all three, both quants' to theirs, and a drafted greedy chat (128 tokens) equals the plain one. The pushed heads differ from the ones CI ran on only in comments and docstrings (the lean check's one-line rule); fmt, lean, the judge's tests, the build andzig build testwere rerun on them.