Skip to content

MLX backend loads full-weight checkpoints as saved: Kev-27B v2 gets an Apple Silicon path (expected ≈ 51 GB, untested at 27B) - #200

Merged
jaredpalmer merged 3 commits into
mainfrom
mlx/full-weight
Sep 30, 2026
Merged

jaredpalmer merged 3 commits into
mainfrom
mlx/full-weight

Conversation

@jaredpalmer

Copy link
Copy Markdown
Owner

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.serve picks the MLX backend on Apple Silicon, and Checkpoint._load_mlx refuses 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_pretrained of transformers' Qwen3_5TextModel, the text backbone alone: tensor names without a prefix (layers.3.self_attn.q_proj.weight), config.json with model_type: qwen3_5_text, no vision tower, no LM head, no MTP. mlx-lm's qwen3_5.Model already accepts a bare text config as its text_config, so kev.mlx_model.load_full builds that model from the checkpoint's own config and loads the shards into it:

  1. language_model.model. is prefixed onto every saved name (mlx-lm nests the backbone there, the same mapping merge_lora makes for peft names).
  2. mlx-lm's own sanitize for a transformers export runs: the DeltaNet conv kernel [C, 1, K] -> [C, K, 1], and the zero-centred RMSNorm weights stored as 1 + w. This is exactly what happens to the Hub bases the LoRA path loads. sanitize decides "this is a transformers export" from the conv kernel's shape, so load_full first checks that every conv kernel really is [C, 1, K]. Otherwise the norms could load unshifted without any error.
  3. Strict: every name and shape must match the model exactly, and every tensor must be in the dtype head.pt names (checked against config.json by the new Checkpoint.saved_dtype, which the torch loader now shares). layer_types must follow full_attention_interval, because that is the only layout mlx-lm builds. Anything else is a ValueError listing what differed. Nothing is ever cast or partially loaded.
  4. The LM head is deleted before anything is materialized. mlx-lm parameters are lazy, so the 2.5 GB head of an untied base like Qwen3.8-27B is never allocated.

Nothing is merged, so the only resident copy is the weights. MLXDecisionModel now takes a built mlx-lm model (load_base + merge_lora for adapters, load_full for 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 own config.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:

Kev-4B (runs/mlx-full-4b) Kev-0.8B (runs/mlx-full-0.8b)
bf16 shards 8.41 GB 1.50 GB
MLX peak at load, full path 8.41 GB 1.50 GB
MLX peak at load, LoRA path (base + merge) 15.87 GB 2.57 GB
full vs LoRA path, max |dp| / flips 0.018 / 0 0.025 / 1
full vs LoRA path with the base's fp32 tensors rounded to bf16 0.0 / 0 (bit-identical) 0.0 / 0
full on MLX vs the same checkpoint on torch fp32 (MPS) 0.015 / 0 0.025 / 0

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_log and the gated-norm weight). The LoRA path keeps them in fp32; a full-weight export holds them in bf16, as every kev.train --full_ft run 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 on mps 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 in runs/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 --remote on 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

  • MoE full-weight checkpoints (qwen3_5_moe_text) are refused. mlx-lm has the module, but nothing here exercises it.
  • /v1/models does not report whether the checkpoint is full or LoRA (lora: 0 already implies it).
  • The Space is untouched: it serves on CUDA, and the checkpoint.py change 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 bf16 equals casting the fp32 export, tensor for tensor; a bf16-trained LoRA is unchanged by the flag.
  • Unit suites (CI job), tests/test_conventions.py, scripts/verify_claims.py (905 claims, including the new README/card numbers that trace to runs/mlx-full-4b/report.json).
  • scripts/mlx_full_parity.py on Kev-4B and Kev-0.8B (table above); kev.serve end to end on the 4B full checkpoint.

jaredpalmer and others added 3 commits September 30, 2026 16:17
…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>
@jaredpalmer

Copy link
Copy Markdown
Owner Author

Thermonuclear self-review (structure, not behaviour):

  • Fixed: the full-weight lora_scale refusal lived in both _full_torch and _load_mlx; it is now one check in Checkpoint.load for either backend.
  • Fixed: _load_mlx read the saved dtype after the mlx-lm import, so a mislabelled checkpoint on a machine without mlx-lm raised ImportError instead of the dtype error. It is read with the other refusals now.
  • Fixed: scripts/mlx_full_parity.py wrote and read JSON by hand; it uses kev.suite.read_json / write_json now.
  • Deduplicated in the PR itself: the config.json / head.pt dtype rule is now Checkpoint.saved_dtype, shared by the torch and MLX loaders, instead of a second copy.
  • Kept on purpose: load_full runs two checks before the strict name/shape/dtype comparison, for off-interval layer_types and non-transformers conv kernels. The name check would catch the first, but the error message would not say why. The second guards mlx-lm's layout heuristic, which would otherwise shift norms silently.
  • Kept: mlx_full_parity.py stays separate from mlx_parity.py. It compares two checkpoints and isolates each load's memory in a subprocess; mlx_parity.py compares one checkpoint's paths plus latency. They share a one-line record filter.
  • Kept: tests/test_mlx.py builds its own tiny tokenizer rather than importing test_unit's fixture. It is a Mac-only module, and its model needs dims mlx-lm's kernels accept (head dims 32).
  • Sizes: kev/mlx_model.py 219 lines, kev/checkpoint.py 356. tests/test_unit.py was already 2.3k lines and gains 30.

Verdict: no structural regressions left that I can see. The new path is one function behind the existing Checkpoint.load dispatch, and MLXDecisionModel got simpler: it takes a built model instead of loading one.

@jaredpalmer
jaredpalmer merged commit d1faa41 into main Sep 30, 2026
2 checks passed
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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant