Repository navigation
MLX backend loads full-weight checkpoints as saved: Kev-27B v2 gets an Apple Silicon path (expected ≈ 51 GB, untested at 27B) - #200
Merged
Conversation
…erge
kev.mlx_model.load_full builds mlx-lm's qwen3_5 model from a full-weight checkpoint's own config.json and shards (strict
names, shapes and dtype; LM head never materialized), so Kev-27B v2 has an Apple Silicon path whose load peak is its
weights. scripts/merge_lora_checkpoint.py --weights_dtype bf16 writes small Kevs as bf16 full weights to check it:
Kev-4B loads at 8.41 GB (= its shards) vs 15.87 GB for the LoRA merge, answers bit-identical to the LoRA path once the
base's fp32 A_log / gated-norm tensors are rounded as a full export holds them, max |dp| 0.015 / 0 flips vs fp32 torch
(runs/mlx-full-{4b,0.8b}). Kev-27B numbers are expected, untested on a Mac.
Generated with [Devin](https://devin.ai)
Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
… backends; read the saved dtype before the mlx import Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Owner
Author
|
Thermonuclear self-review (structure, not behaviour):
Verdict: no structural regressions left that I can see. The new path is one function behind the existing |
4 tasks done
jaredpalmer
added a commit
that referenced
this pull request
Oct 1, 2026
…loading #200, prefix cache make_room + out-of-memory retry fix #201, fp32 long-row attention #202) The skills pinned 2ea5660. 71d4829 is main after #200, #201 and #202: the server's prefix cache now drops the entries a batch will evict before its pass (PrefixCache.make_room), so a new near-limit state no longer sits next to the old one, and the out-of-memory retry also counts this batch's own hits as held states. Checked on an H200 deploy of Kev-27B (app kev-ref71d4829-27b, stopped; runs/kev-deploy-71d4829/deploy_check.json): semif-v1 development over --remote vs runs/rel27-public max |dp| 0.0122, 0 flips (252 rows); the 70,065-token state gets the 422; two ~60k-token states answered (8.1-10.1 s new, 0.28 s repeated), the second evicting the first cleanly (cached_states 1, oom_retries 0). SKILL.md's long-document latencies updated. Spend <= ~$0.45. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
jaredpalmer
added a commit
that referenced
this pull request
Oct 1, 2026
…loading #200, prefix cache make_room + out-of-memory retry fix #201, fp32 long-row attention #202) (#204) The skills pinned 2ea5660. 71d4829 is main after #200, #201 and #202: the server's prefix cache now drops the entries a batch will evict before its pass (PrefixCache.make_room), so a new near-limit state no longer sits next to the old one, and the out-of-memory retry also counts this batch's own hits as held states. Checked on an H200 deploy of Kev-27B (app kev-ref71d4829-27b, stopped; runs/kev-deploy-71d4829/deploy_check.json): semif-v1 development over --remote vs runs/rel27-public max |dp| 0.0122, 0 flips (252 rows); the 70,065-token state gets the 422; two ~60k-token states answered (8.1-10.1 s new, 0.28 s repeated), the second evicting the first cleanly (cached_states 1, oom_retries 0). SKILL.md's long-document latencies updated. Spend <= ~$0.45. Generated with [Devin](https://devin.ai) Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
seanpaulconnelly
added a commit
to seanpaulconnelly/kev
that referenced
this pull request
Oct 3, 2026
…for v2 (full weights) and v1 (LoRA), and v1 parity v2 (jaredpalmer/kev-27b @ af0e6d5, full bf16 weights, loaded by the MLX path from jaredpalmer#200), served on an M5 Max with 128 GB: transfer-v4 development 0.849 against the published 0.851, Brier 0.218 against 0.218 (runs/mlx-27b-v2-transfer-v4). Server footprint 53.7 GB steady, 54.4 GB peak, for 51.2 GB of weights. v1 (tag v1-lora), served through the adapter-merge path: 0.849 against 0.848 (runs/mlx-27b-transfer-v4); max |dp| 0.022 against the torch bf16 path, 0 flips on 40 decision-v7 development questions (runs/mlx-parity-27b).
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
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.
Kev-27B v2 is a full-weight checkpoint: 51 GB of bf16 backbone weights (
config.json+ 11 safetensors shards) and a pointer head, with no adapter. On a Mac it has no path.kev.servepicks the MLX backend on Apple Silicon, andCheckpoint._load_mlxrefuses full weights ("the MLX backend merges an adapter into the base"). The torch fallback on MPS runs the reference DeltaNet code, which is far too slow for a 27B. #175 (Sean Connelly) showed v1, the LoRA version, serving through MLX on a 128 GB M5 Max: transfer-v4 development 0.849 against the published 0.848, 52 GB steady and 97 GB at peak while the adapter was merged into the base. v2 needs its own way in.Mechanism
A full-weight checkpoint is
save_pretrainedof transformers'Qwen3_5TextModel, the text backbone alone: tensor names without a prefix (layers.3.self_attn.q_proj.weight),config.jsonwithmodel_type: qwen3_5_text, no vision tower, no LM head, no MTP. mlx-lm'sqwen3_5.Modelalready accepts a bare text config as itstext_config, sokev.mlx_model.load_fullbuilds that model from the checkpoint's own config and loads the shards into it:language_model.model.is prefixed onto every saved name (mlx-lm nests the backbone there, the same mappingmerge_loramakes for peft names).sanitizefor a transformers export runs: the DeltaNet conv kernel[C, 1, K] -> [C, K, 1], and the zero-centred RMSNorm weights stored as1 + w. This is exactly what happens to the Hub bases the LoRA path loads.sanitizedecides "this is a transformers export" from the conv kernel's shape, soload_fullfirst checks that every conv kernel really is[C, 1, K]. Otherwise the norms could load unshifted without any error.config.jsonby the newCheckpoint.saved_dtype, which the torch loader now shares).layer_typesmust followfull_attention_interval, because that is the only layout mlx-lm builds. Anything else is aValueErrorlisting what differed. Nothing is ever cast or partially loaded.Nothing is merged, so the only resident copy is the weights.
MLXDecisionModelnow takes a built mlx-lm model (load_base+merge_lorafor adapters,load_fullfor full weights); everything downstream (encoder, rows, prompt-cache prefix, pointer head) is unchanged.backend="auto"now picks MLX for full weights too.hybrid_base()reads a full checkpoint's ownconfig.json, because that is the backbone that actually loads. The refusals stay as they were (lora_scale, option_isolation, attention-only backbones), all before the mlx-lm import.How it was checked without a 27B-capable Mac
This Mac has 32 GB, so I checked the path on the small Kevs written as full weights.
scripts/merge_lora_checkpoint.py --weights_dtype bf16(new) merges an fp32-trained LoRA in fp32 and rounds each tensor to bf16 once: round(fp32(W) + fp32(delta)), the arithmetic of the served merge.scripts/mlx_full_parity.py(new) then runs each load in its own process, so every memory peak belongs to that load alone. Numbers are from 60 (4B) / 100 (0.8B) clean decision-v7 development questions on an M5 32 GB:runs/mlx-full-4b)runs/mlx-full-0.8b)One thing surprised me. I expected full vs LoRA to agree to about 1e-3, but the gap is the size of bf16 noise. The cause is that the Qwen3.5 bases store two small tensors per DeltaNet layer in fp32 (
A_logand the gated-norm weight). The LoRA path keeps them in fp32; a full-weight export holds them in bf16, as everykev.train --full_ftrun trains and saves them (Kev-27B's index is exactly 2 bytes per parameter). Round just those tensors in the LoRA model and the two paths give bit-identical probabilities. So the merge arithmetic and the layout mapping are exact, and the remaining difference is a property of the checkpoint, not of the loader. The full-vs-torch row (0.015, 0 flips) is in line with the committed LoRA parity runs (runs/mlx-parity-4b: 0.024, 0 flips).kev.serve --run <4B full checkpoint>ends up onmps via mlx (bfloat16). The README request returns 200 (515 ms with a new state, 151 ms with the state cached), and a 70,002-token state gets the usual 422 (recorded inruns/mlx-full-4b/report.json["serve"]).Kev-27B: expected, not measured
The 4B and 0.8B loads peak at exactly their shard size, so v2 should need about 51.2 GB plus working memory, against 97 GB for v1's merge. The README, the Kev-27B card and AGENTS.md now say: a 64 GB Mac is borderline (51 GB is above macOS's default GPU working-set recommendation of about three quarters of RAM), and 96–128 GB should fit. They say it is expected from small-model measurements and not yet measured on a large Mac, and they credit #175 for the v1 result. If someone with a 128 GB Mac runs
kev.benchmark --remoteon transfer-v4 development plus the footprint, that settles it. The decision rule written before any run: accuracy within ±0.015 of the published 0.851 and Brier within ±0.02, as #175 used.Not in this PR
qwen3_5_moe_text) are refused. mlx-lm has the module, but nothing here exercises it./v1/modelsdoes not report whether the checkpoint is full or LoRA (lora: 0already implies it).checkpoint.pychange only affects backend resolution for MLX.Test plan
tests/test_mlx.py(Apple Silicon, 39 s, 10 passed): a random 4-layer hybrid full-weight checkpoint loads on MLX with every weight equal to the saved tensor after the layout rules, no LM head, hidden states within bf16 noise of torch fp32, and within the Metal fp32 matmul's precision in fp32 (0.0036 relative). Refusals: a missing, unexpected, wrong-shape or wrong-dtype tensor, MLX-layout conv kernels, another model_type, off-interval layer_types, a mislabelled dtype, option_isolation, lora_scale, an attention-only backbone. The pinned Kev-0.8B merged to bf16 full weights is bit-identical to the LoRA path (fp32 tensors rounded) and loads within 1.01× its weight bytes.tests/test_unit.py::test_merged_lora_checkpoint_in_bf16_rounds_the_fp32_merge_once(CPU):--weights_dtype bf16equals casting the fp32 export, tensor for tensor; a bf16-trained LoRA is unchanged by the flag.tests/test_conventions.py,scripts/verify_claims.py(905 claims, including the new README/card numbers that trace toruns/mlx-full-4b/report.json).scripts/mlx_full_parity.pyon Kev-4B and Kev-0.8B (table above);kev.serveend to end on the 4B full checkpoint.