Repository navigation
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
Draft
jaredpalmer wants to merge 2 commits into
jaredpalmer wants to merge 2 commits into
Conversation
…_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>
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 treats a checkpoint as data: a pointer head, weights, and a tokenizer. On Apple Silicon one loader path can break that.
kev.serveloads a hybrid checkpoint throughkev.mlx_model.load_base, and mlx-lm 0.31.3, the version pinned inuv.lock, imports and runs the Python file a base'sconfig.jsonnames undermodel_file. The base ismeta.basefrom the checkpoint's ownhead.pt, so the checkpoint chooses which config is read.Pathjoin treats an absolutemodel_fileas that path, so the file does not have to live inside the snapshot. A freshresolve_rundownload still excludes*.py. The torch backend never does this. transformers is called withouttrust_remote_code.load_basereadsconfig.jsonand raisesValueErrornaming the field beforeload_modelruns. Kev's bases are mlx-lm's built-inqwen3_5architectures, so a published base loads as before, with one extra JSON read.read_metapassesweights_only=True. On torch 2.8 the default is already weights-only: ahead.ptwhose pickle would create a file does not run, with or without the argument. The explicit argument is for the Space, which vendorskev/checkpoint.pyonto ZeroGPU's torch, and nothing in this repo pins that torch.jaredpalmer/kev-0.8bandjaredpalmer/kev-4bhead.ptboth load. The temperatures come back 2.35 and 2.41.This keeps Wenjix's commit from #262. I dropped the
AGENTS.mdlayout clause. The refusal belongs onload_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_onlyreaches the live Space only whenscripts/publish_space.shruns. mlx-lm 0.32.2 added its owntrust_remote_codegate formodel_file. We still pinmlx-lm<0.32, and this check stays correct after a bump.Test plan
tests/test_unit.py::test_read_meta_refuses_a_pickle_payloadon torch 2.8. The sentinel file is not created.read_metaon the publishedjaredpalmer/kev-0.8bandjaredpalmer/kev-4bhead.ptfiles withweights_only=True.test_resume_is_bit_identicalfor1-gate0and1-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_codeneeds Apple Silicon and mlx-lm. The guard matches mlx-lm 0.31.3load_model:config.get("model_file") is not None, thenexec_module. I could not run the sentinel here.