Skip to content

Loading a checkpoint never runs code it brings: refuse a base's model_file on MLX, read head.pt weights-only - #274

Draft
jaredpalmer wants to merge 2 commits into
mainfrom
cursor/mlx-refuse-model-file-519a
Draft

jaredpalmer wants to merge 2 commits into
mainfrom
cursor/mlx-refuse-model-file-519a

Conversation

@jaredpalmer

Copy link
Copy Markdown
Owner

Kev treats a checkpoint as data: a pointer head, weights, and a tokenizer. On Apple Silicon one loader path can break that. kev.serve loads a hybrid checkpoint through kev.mlx_model.load_base, and mlx-lm 0.31.3, the version pinned in uv.lock, imports and runs the Python file a base's config.json names under model_file. The base is meta.base from the checkpoint's own head.pt, so the checkpoint chooses which config is read. Path join treats an absolute model_file as that path, so the file does not have to live inside the snapshot. A fresh resolve_run download still excludes *.py. The torch backend never does this. transformers is called without trust_remote_code.

load_base reads config.json and raises ValueError naming the field before load_model runs. Kev's bases are mlx-lm's built-in qwen3_5 architectures, so a published base loads as before, with one extra JSON read. read_meta passes weights_only=True. On torch 2.8 the default is already weights-only: a head.pt whose pickle would create a file does not run, with or without the argument. The explicit argument is for the Space, which vendors kev/checkpoint.py onto ZeroGPU's torch, and nothing in this repo pins that torch. jaredpalmer/kev-0.8b and jaredpalmer/kev-4b head.pt both load. The temperatures come back 2.35 and 2.41.

This keeps Wenjix's commit from #262. I dropped the AGENTS.md layout clause. The refusal belongs on load_base, next to the call that would have executed the file. The module index does not need another sentence.

Not in this PR: republishing the Space. weights_only reaches the live Space only when scripts/publish_space.sh runs. mlx-lm 0.32.2 added its own trust_remote_code gate for model_file. We still pin mlx-lm<0.32, and this check stays correct after a bump.

Test plan

  • tests/test_unit.py::test_read_meta_refuses_a_pickle_payload on torch 2.8. The sentinel file is not created.
  • read_meta on the published jaredpalmer/kev-0.8b and jaredpalmer/kev-4b head.pt files with weights_only=True.
  • test_resume_is_bit_identical for 1-gate0 and 1-gate4. Both pass when run alone. They failed earlier only while another suite was on the CPU.
  • tests/test_mlx.py::test_load_base_refuses_model_code needs Apple Silicon and mlx-lm. The guard matches mlx-lm 0.31.3 load_model: config.get("model_file") is not None, then exec_module. I could not run the sentinel here.
  • CI
Open in Web Open in Cursor 

Wenjix and others added 2 commits October 27, 2026 13:44
…_file on MLX, read head.pt weights-only

mlx-lm 0.31's load_model imports and runs the Python module a base's
config.json names under model_file, and the base is named by the
checkpoint's own head.pt. load_base now refuses such a base before
mlx-lm sees it. read_meta passes weights_only=True explicitly, so
head.pt is unpickled weights-only on torch versions older than 2.6
too (the Space vendors kev/checkpoint.py onto an unpinned torch).

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
The module index already describes what mlx_model is. The refusal lives on load_base, next to the call that would have executed the file.

Co-authored-by: Jared Palmer <hello@jaredpalmer.com>
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.

3 participants