xmd14/tensorizer-torch-compat-rce-poc

PoC for Huntr Report — tensorizer_loading Silently Overrides PyTorch's Safe weights_only Default, Enabling Arbitrary Code Execution

Target: Tensorizer (CoreWeave) — github.com/coreweave/tensorizer

CWE: CWE-502 (Deserialization of Untrusted Data), secondary CWE-693 (Protection Mechanism Failure)

Affected versions (both verified): tensorizer==2.12.1 (latest PyPI release) and current main branch on GitHub (fetched and compared directly; identical relevant code).

Tested against: torch==2.13.0+cpu (latest stable, live-tested) and torch==2.6.0 (the release that introduced the weights_only=True default; verified by reading the tagged v2.6.0 source directly -- the relevant _default_to_weights_only function and the weights_only/pickle_module resolution logic are byte-identical to 2.13.0).

Important note on the PoC file

evil_checkpoint.pt is an ordinary PyTorch pickle checkpoint (created with plain pickle.dump), not Tensorizer's native .tensors binary format. It represents exactly the kind of file a developer would pass to torch.load() -- which is the documented use case for Tensorizer's torch_compat compatibility layer. Tensorizer's native .tensors format used directly via TensorSerializer/TensorDeserializer does not use pickle and is not implicated by this report.

The payload is intentionally harmless: it only creates a local marker file (tensorizer_rce_marker.txt) to prove code execution. It does not delete files, exfiltrate data, or download anything.

1. The documented, general scenario

Tensorizer ships a documented compatibility layer, tensorizer.torch_compat, intended as a drop-in speedup for existing, unmodified torch.save/torch.load call sites:

from tensorizer.torch_compat import tensorizer_loading

with tensorizer_loading(device="cuda", num_readers=4):
    deserialized_module = torch.load("module.pt")

This is CoreWeave's own usage example (from the module's docstring). The whole point of the API is that you do not need to change your existing torch.load(...) call -- you just wrap it in this context manager to get tensorizer's faster loading.

Tensorizer is explicitly designed to load models from local storage, HTTP/HTTPS, or S3 endpoints -- i.e. from sources that may not be fully trusted.

2. Baseline: the same file is safely rejected outside the context

PyTorch 2.6+ defaults torch.load() to weights_only=True whenever no custom pickle_module is supplied, which blocks arbitrary globals in untrusted checkpoints:

[1. torch.load(evil_checkpoint.pt) OUTSIDE tensorizer_loading(), default weights_only] -> blocked
    exception: UnpicklingError: Weights only load failed. In PyTorch 2.6, we changed
    the default value of the `weights_only` argument in `torch.load` from `False` to
    `True`. ... it can result in arbitrary code execution ...

3. The same file, same call, executes arbitrary code inside the documented context

with tensorizer_loading():
    torch.load("evil_checkpoint.pt")   # identical call, no weights_only argument

[2. torch.load(evil_checkpoint.pt) INSIDE tensorizer_loading(), default weights_only] -> CODE EXECUTED
    exception: RuntimeError: Invalid magic number; corrupt file?
    warnings captured during case 2: 0
      (none -- the safe-default override is completely silent)

The arbitrary command runs (verified via the tensorizer_rce_marker.txt marker file created by the payload) even though torch.load() subsequently raises an unrelated error while trying to interpret the returned object as checkpoint data. No warning of any kind is emitted about the safe default being overridden -- confirmed by capturing all Python warnings around the call.

The only difference between the safe case (2) and the vulnerable case (3) is being inside tensorizer_loading().

As a control, explicitly passing weights_only=True inside the same context is correctly rejected (the user can re-enable the protection explicitly):

[3. torch.load(evil_checkpoint.pt, weights_only=True) INSIDE tensorizer_loading()] -> blocked
    exception: RuntimeError: Can not safely load weights when explicit pickle_module is specified

And, as a sanity baseline showing this is the same underlying danger class as ordinary insecure loading (just triggered implicitly instead of explicitly), weights_only=False outside the context also executes:

[4. torch.load(evil_checkpoint.pt, weights_only=False) OUTSIDE tensorizer_loading()] -> CODE EXECUTED
    exception: RuntimeError: Invalid magic number; corrupt file?

4. Root cause

tensorizer/torch_compat.py's _load_wrapper (the function that replaces torch.load while tensorizer_loading() is active) unconditionally forces a custom pickle_module:

return _ORIG_TORCH_LOAD(
    f, map_location, pickle_module=_tensorizer_pickle, *args,
    weights_only=weights_only, **kwargs,
)

In torch/serialization.py (PyTorch, unchanged from v2.6.0 through the current 2.13.0):

def _default_to_weights_only(pickle_module):
    is_fbcode = not hasattr(torch.version, "git_version")
    return pickle_module is None and not is_fbcode

weights_only_not_set = weights_only is None
if weights_only_not_set:
    weights_only = _default_to_weights_only(pickle_module)
...
if weights_only:
    if pickle_module is not None:
        raise RuntimeError(
            "Can not safely load weights when explicit pickle_module is specified"
        )

PyTorch only selects the safe weights_only=True default when pickle_module is None. Because tensorizer_loading() injects a non-None pickle_module into every torch.load() call made while it is active -- including calls the developer never modified and that would otherwise get the safe default -- this silently changes the effective default to unrestricted pickle loading, unless the caller explicitly passes weights_only=True. The explicit RuntimeError guard at the bottom only fires when weights_only is already True by the time it is checked; with an unspecified (None) value, resolution to False happens first, so that guard is never reached.

5. Supplementary internal evidence (not the primary finding)

tensorizer.torch_compat._TensorizerUnpickler(pickle.Unpickler) restricts persistent_load (the mechanism used to reference tensor storages), but never overrides find_class anywhere in the module (verified by exhaustive search of the package source). find_class is the actual gateway used by ordinary REDUCE/GLOBAL pickle opcodes. This corroborates why case 2 above executes code, but is not by itself the primary finding, since this class is internal and not meant to be used directly by application code -- see internal_find_class_evidence.py.

Impact

PyTorch 2.6 and later default torch.load() to weights_only=True when no custom pickle module is supplied, preventing arbitrary globals in untrusted checkpoints from executing.

Tensorizer's documented tensorizer_loading() compatibility context automatically injects its custom pickle module into ordinary torch.load() calls. Because PyTorch only selects the safe weights-only default when pickle_module is None, this silently changes the effective default to unrestricted pickle loading unless the caller explicitly passes weights_only=True.

Using the same malicious checkpoint and the same torch.load() call:

• Outside tensorizer_loading(): blocked by PyTorch's default weights-only loader.
• Inside tensorizer_loading(): arbitrary system command executed.
• Inside tensorizer_loading() with weights_only=True explicitly: blocked.
• Outside the context with weights_only=False explicitly: arbitrary command executed as expected.

No warning is emitted when the safe default is overridden.

This affects users who follow Tensorizer's documented compatibility workflow and load PyTorch checkpoints from model repositories, object storage, shared projects, or other sources they do not fully trust. The issue is in Tensorizer's torch compatibility layer and does not imply that Tensorizer's native .tensors format uses pickle.

Environment

tensorizer==2.12.1
torch==2.13.0+cpu
Python 3.13.3
Windows 11 Home 10.0.26200 (Build 26200)

Files in this repository

File	Purpose
evil_checkpoint.pt	The PoC malicious checkpoint (ordinary PyTorch pickle, not .tensors format). Loading it unsafely creates tensorizer_rce_marker.txt -- a harmless proof of execution, nothing destructive.
generate_evil_checkpoint.py	Regenerates the PoC checkpoint.
test_matrix.py	The primary, self-contained PoC. Runs all 4 cases above plus the warnings-capture check.
raw_output.txt	Unedited console output from running test_matrix.py.
internal_find_class_evidence.py	Supplementary evidence only (section 5 above), not the primary PoC.

Reproduction

pip install tensorizer torch
python test_matrix.py

Suggested fixes

• Do not automatically inject a custom pickle_module into torch.load() calls for files that are not confirmed to be Tensorizer-produced sidecar-paired files.
• Explicitly preserve weights_only=True behavior when the caller does not specify weights_only, rather than letting it silently resolve to False as a side effect of the custom pickle_module.
• Require an explicit opt-in (e.g. allow_unsafe_pickle=True) before falling back to unrestricted pickle loading.
• Emit a clear security warning whenever the safe default is overridden as a side effect of entering tensorizer_loading().
• Add a regression test asserting that a malicious checkpoint is rejected both inside and outside tensorizer_loading() unless unsafe loading is explicitly requested.
• Document plainly (not just in the suppress_weights_only parameter's docstring) that weights_only=False must never be used, implicitly or explicitly, with untrusted checkpoint files.