From 3c0372532b8ad41f8e5e2a8012433af27a64a292 Mon Sep 17 00:00:00 2001 From: Wenjix Date: Tue, 27 Oct 2026 02:15:04 -0700 Subject: [PATCH 1/2] Loading a checkpoint never runs code it brings: refuse a base's model_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> --- AGENTS.md | 2 +- kev/checkpoint.py | 2 +- kev/mlx_model.py | 9 ++++++++- tests/test_mlx.py | 14 +++++++++++++- tests/test_unit.py | 16 ++++++++++++++++ 5 files changed, 39 insertions(+), 4 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index deae30dc..ea063875 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -385,7 +385,7 @@ runs / the endpoint / the volumes. Tests: `tests/test_skill_scripts.py`. `encode_batch` + `batch_loss` (CE, anchor KL, permutation KL, RPS, smoothing/Brier/focal), `main` orchestration. Grad accumulation over small padded batches. - `kev/anchors.py` frozen-base zero-shot distributions per training question, the target for `--anchor_w` - `kev/checkpoint.py` Checkpoint (resolve run dir or Hub id, `head.pt` schema = `Meta`, the LoRA-vs-full loader rule `full`/`shards`/`weights_sha256` (a full checkpoint loads in its head.pt `weights_dtype`, which must match config.json's `dtype`), load with `LoadOptions` -> torch `DecisionModel` or, with `backend="mlx"`/`"auto"`, `MLXDecisionModel`; `warm_start` for deltas) -- `kev/mlx_model.py` Apple Silicon backend: mlx-lm Qwen3.5 backbone, LoRA merged in fp32 on the CPU stream (`merge_lora`) or a full-weight checkpoint loaded as saved (`load_full`), Kev's encoder/rows and the torch PointerHead unchanged; state prefix = mlx-lm prompt cache filled `PREFILL_CHUNK` tokens per pass, branches on a replicated copy +- `kev/mlx_model.py` Apple Silicon backend: mlx-lm Qwen3.5 backbone, LoRA merged in fp32 on the CPU stream (`merge_lora`) or a full-weight checkpoint loaded as saved (`load_full`), Kev's encoder/rows and the torch PointerHead unchanged; state prefix = mlx-lm prompt cache filled `PREFILL_CHUNK` tokens per pass, branches on a replicated copy; `load_base` refuses a base whose config.json names a `model_file` (mlx-lm 0.31.x would execute it unconditionally) - `kev/device.py` default device and explicit selection (`select`, a child-process probe: CUDA/ROCm, then MPS, then CPU; `KEV_DEVICE_PROBE_TIMEOUT`, default 30s) / sync / empty_cache / allocated_bytes for cuda, mps, cpu - `kev/metrics.py` pure-numpy scoring of benchmark rows: ECE, Brier, NLL, selective prediction (tie-aware coverage@error, AURC), temperature fit, out-of-fold CV calibration report (`cross_validated_temperature`), paired bootstrap - `kev/predictors.py` LocalPredictor (checkpoint), RemotePredictor (System One endpoint), JevPredictor (AI SDK worker) diff --git a/kev/checkpoint.py b/kev/checkpoint.py index 27246247..16aa011f 100644 --- a/kev/checkpoint.py +++ b/kev/checkpoint.py @@ -74,7 +74,7 @@ def to_dict(self): def read_meta(run): - return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu")) + return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu", weights_only=True)) def write_meta(run, meta): diff --git a/kev/mlx_model.py b/kev/mlx_model.py index 0042e36b..7a4dd29b 100644 --- a/kev/mlx_model.py +++ b/kev/mlx_model.py @@ -37,7 +37,14 @@ def load_base(base_dir): - """mlx-lm's model of a Hub base snapshot, weights as stored (bf16 for the Qwen3.5 bases); the LoRA path merges into it.""" + """mlx-lm's model of a Hub base snapshot, weights as stored (bf16 for the Qwen3.5 bases); the LoRA path merges into it. + A config.json that names a `model_file` is refused: mlx-lm 0.31's load_model imports and runs that module to build + the model (0.32.2 asks for trust_remote_code first), and the base is whatever the checkpoint's head.pt names, so a + checkpoint could choose code to run. Kev's bases are mlx-lm's built-in architectures; the torch backend ignores the field.""" + config = json.loads((Path(base_dir) / "config.json").read_text(encoding="utf-8")) + if config.get("model_file") is not None: + raise ValueError(f"{base_dir}: config.json names model_file {config['model_file']!r}, model code mlx-lm would import and run; " + "the MLX backend loads only mlx-lm's built-in architectures. Load this checkpoint with backend=torch (KEV_BACKEND=torch), which ignores the field") mx.set_cache_limit(CACHE_LIMIT) return load_model(Path(base_dir))[0] diff --git a/tests/test_mlx.py b/tests/test_mlx.py index 46870a30..2ec3902a 100644 --- a/tests/test_mlx.py +++ b/tests/test_mlx.py @@ -15,7 +15,7 @@ pytestmark = pytest.mark.skipif(platform.system() != "Darwin" or platform.machine() != "arm64", reason="MLX runs on Apple Silicon only") from kev.checkpoint import Checkpoint, LoadOptions, mlx_available # noqa: E402 -from kev.mlx_model import MLXDecisionModel, merge_lora # noqa: E402 +from kev.mlx_model import MLXDecisionModel, load_base, merge_lora # noqa: E402 from kev.model import SCORING_INTERFACE, DecisionModel # noqa: E402 RUN = "jaredpalmer/kev-0.8b@9a45d25eb2ab761841196625383fa1dff0e56c1e" # pinned (round 15, 2026-09-24): bf16 noise is per checkpoint, so a republish must not move these bars @@ -65,6 +65,18 @@ def test_backend_resolution(): Checkpoint("jaredpalmer/kev-4b@qwen3").load("mps", LoadOptions(backend="mlx")) # Qwen3 base: no DeltaNet layers +def test_load_base_refuses_model_code(tmp_path): + """mlx-lm 0.31.x executes config.json's model_file unconditionally; a base that brings one is refused before + load_model gets as far as running it (valid safetensors, so only the guard stands in the way).""" + sentinel = tmp_path / "sentinel" + (tmp_path / "config.json").write_text(json.dumps({"model_type": "qwen3_5", "model_file": "model.py"}), encoding="utf-8") + mx.save_safetensors(str(tmp_path / "model.safetensors"), {"x": mx.zeros((1,))}) + (tmp_path / "model.py").write_text(f"open({str(sentinel)!r}, 'w').write('x')\n", encoding="utf-8") + with pytest.raises(ValueError, match="model_file"): + load_base(tmp_path) + assert not sentinel.exists() + + # --- full-weight checkpoints (kev.train --full_ft; kev.mlx_model.load_full): a 4-layer Qwen3.5 text model with random # weights (norms, conv kernels and A_log included, so a missed layout rule shows), saved as a full-weight run saves it diff --git a/tests/test_unit.py b/tests/test_unit.py index 0557c4df..4b3f74b4 100644 --- a/tests/test_unit.py +++ b/tests/test_unit.py @@ -185,6 +185,22 @@ def test_checkpoint_meta_round_trip_and_defaults(tmp_path): LoadOptions.from_env({"KEV_BACKEND": "metal"}) +def test_read_meta_refuses_a_pickle_payload(tmp_path): + """read_meta loads head.pt with weights_only=True: a checkpoint whose head.pt unpickles into code never runs it + (the Space vendors checkpoint.py onto ZeroGPU's unpinned torch, where the default was not always safe).""" + import pickle + from kev.checkpoint import Meta, read_meta, write_meta + sentinel = tmp_path / "sentinel" + class Payload: + def __reduce__(self): + return (Path.touch, (sentinel,)) + meta = Meta.from_dict({"head": {"w": torch.zeros(1)}, "base": "b"}); meta.extra["payload"] = Payload() + write_meta(tmp_path, meta) + with pytest.raises(pickle.UnpicklingError): + read_meta(tmp_path) + assert not sentinel.exists() + + def test_fused_default_needs_pinned_fla(monkeypatch): """kev.serve's CUDA fused default (checkpoint.fused_available): on only with flash-linear-attention importable at fused_qwen35.FLA_VERSION; it is not in the serve extra, so a plain install serves unfused instead of failing to import fla.""" From 9d39436654cc692620789a0e5ebec747aabe2fa9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 27 Oct 2026 13:45:05 +0000 Subject: [PATCH 2/2] Drop the AGENTS.md layout note for the model_file refusal. 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 --- AGENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index ea063875..deae30dc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -385,7 +385,7 @@ runs / the endpoint / the volumes. Tests: `tests/test_skill_scripts.py`. `encode_batch` + `batch_loss` (CE, anchor KL, permutation KL, RPS, smoothing/Brier/focal), `main` orchestration. Grad accumulation over small padded batches. - `kev/anchors.py` frozen-base zero-shot distributions per training question, the target for `--anchor_w` - `kev/checkpoint.py` Checkpoint (resolve run dir or Hub id, `head.pt` schema = `Meta`, the LoRA-vs-full loader rule `full`/`shards`/`weights_sha256` (a full checkpoint loads in its head.pt `weights_dtype`, which must match config.json's `dtype`), load with `LoadOptions` -> torch `DecisionModel` or, with `backend="mlx"`/`"auto"`, `MLXDecisionModel`; `warm_start` for deltas) -- `kev/mlx_model.py` Apple Silicon backend: mlx-lm Qwen3.5 backbone, LoRA merged in fp32 on the CPU stream (`merge_lora`) or a full-weight checkpoint loaded as saved (`load_full`), Kev's encoder/rows and the torch PointerHead unchanged; state prefix = mlx-lm prompt cache filled `PREFILL_CHUNK` tokens per pass, branches on a replicated copy; `load_base` refuses a base whose config.json names a `model_file` (mlx-lm 0.31.x would execute it unconditionally) +- `kev/mlx_model.py` Apple Silicon backend: mlx-lm Qwen3.5 backbone, LoRA merged in fp32 on the CPU stream (`merge_lora`) or a full-weight checkpoint loaded as saved (`load_full`), Kev's encoder/rows and the torch PointerHead unchanged; state prefix = mlx-lm prompt cache filled `PREFILL_CHUNK` tokens per pass, branches on a replicated copy - `kev/device.py` default device and explicit selection (`select`, a child-process probe: CUDA/ROCm, then MPS, then CPU; `KEV_DEVICE_PROBE_TIMEOUT`, default 30s) / sync / empty_cache / allocated_bytes for cuda, mps, cpu - `kev/metrics.py` pure-numpy scoring of benchmark rows: ECE, Brier, NLL, selective prediction (tie-aware coverage@error, AURC), temperature fit, out-of-fold CV calibration report (`cross_validated_temperature`), paired bootstrap - `kev/predictors.py` LocalPredictor (checkpoint), RemotePredictor (System One endpoint), JevPredictor (AI SDK worker)