wildbits/nemo-mfv-poc

NVIDIA NeMo Model File Vulnerability PoC

This repository contains a harmless proof-of-concept for a NVIDIA NeMo .nemo model restoration issue.

Summary

A crafted .nemo archive can include an attacker-controlled model_weights.ckpt file. During model restoration, NeMo loads this checkpoint using PyTorch deserialization.

The PoC demonstrates that attacker-controlled pickle content can execute during the public restore_from() loading workflow.

Safety

The payload is intentionally harmless. It only creates this marker file:

/tmp/nemo_mfv_public_restore_dict_marker.txt

No destructive action is performed.

Reproduction

Create a clean Python environment and install dependencies:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -U pip
pip install -r requirements.txt

Run:

python reproduce.py

Expected PoC result:

[+] Marker after restore: True
[+] Marker content:
NeMo public restore dict payload marker

The restore process may raise an exception about an unexpected key in the state_dict. This happens after the payload has already executed, which demonstrates that unsafe deserialization occurs before state_dict validation.

Affected Component

• Framework: NVIDIA NeMo
• Format: .nemo
• File inside archive: model_weights.ckpt
• Loading API: restore_from()
• Sink: torch.load(..., weights_only=False)

Impact

If a user or automated ML pipeline restores a malicious .nemo model from an untrusted source, attacker-controlled code may execute during model loading.