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:
1[+] Marker after restore: True
2[+] Marker content:
3NeMo 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.
