NVIDIA-Merlin / NVIDIA-Merlin/Transformers4Rec

[SECURITY] `torch.load()` called without `weights_only=True` under `torch>=1.0` pin

Open
#815 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.3k
Forks
165
Avg merge
1m
Merged PRs (30d)
2

Description

Bug description

transformers4rec.torch.model.base.Model.load (rewritten in PR #808, commit 41b14d7b) and its accompanying test call torch.load(...) without passing weights_only=True:

requirements/pytorch.txt currently pins only torch>=1.0:

https://github.com/NVIDIA-Merlin/Transformers4Rec/blob/8bf122f5dcb39feecfc6dabde734d79c2d1c4380/requirements/pytorch.txt

On installs of PyTorch <2.6 (which the pin permits), torch.load defaults to weights_only=False. That is unrestricted pickle deserialization: any __reduce__ on the file executes on load, which is the same class of issue as CVE-2025-32434. PyTorch ≥2.6 flipped the default to weights_only=True, but T4Rec's pin does not enforce that.

This is particularly relevant because PRs #802/#804 were framed as pickle-security hardening — the hardening does not cover this code path. When the allowlist-based utils/serialization.py was removed in PR #808, torch.load became the primary load path again, with its version-dependent default.

Steps/Code to reproduce bug

Benign repro (side effect is a temp file; no persistent change):

import os, tempfile, torch

class Exploit:
    def __reduce__(self):
        return (os.system, ("touch /tmp/t4rec_rce_demo",))

with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tf:
    torch.save(Exploit(), tf.name)
    path = tf.name

torch.load(path, weights_only=False)    # runs os.system under pickle's __reduce__
assert os.path.exists("/tmp/t4rec_rce_demo")

torch.load(path, weights_only=True) refuses the same payload with UnpicklingError:

Weights only load failed. [...] Trying to load unsupported GLOBAL posix.system whose module posix is blocked.

Expected behavior
  • Every internal torch.load call passes weights_only=True.
  • Tutorials and docstrings are updated to use the same pattern.
  • requirements/pytorch.txt bumps the floor to torch>=2.6 so the safe default is always in effect.
Environment details
  • Transformers4Rec: main @ 8bf122f5
  • PyTorch: the pin is >=1.0 (see requirements/pytorch.txt); installs on <2.6 are the affected population.
  • Python: any
Additional context

Two-part remediation:

  1. In code: add weights_only=True to every torch.load(...) call in the library, tests, and example notebooks. This is safe because the surrounding code already expects a state_dict (plain tensor dict), which weights_only=True supports.

  2. In requirements: bump torch>=1.0 to torch>=2.6. The project's other dependencies (HF transformers, Merlin) already require a recent PyTorch, so this bump is effectively a no-op for supported users.

Happy to send a PR covering both parts.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with transformers4rec.torch.model.base.Model.load, its docstring, and the call in tests/unit/torch/model/test_model.py around line 425. Search the library, tests, and example notebooks for other torch.load calls, then inspect requirements/pytorch.txt. Done means every load uses weights_only=True, guidance matches, and the relevant unit tests pass with the PyTorch floor set to 2.6.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, security
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.