The Trouble with Tensors: Securing Tensor Serialization
Saving PyTorch tensors is a routine part of machine learning, from checkpointing mid-training to publishing final model weights on Hugging Face. But for years, the standard way to load those weights exposed anyone who did it to remote code execution. PyTorch 2.6 finally changed the default, and if you hit an UnpicklingError after upgrading, this talk explains exactly why that error exists. We go deep on PyTorch's Pickle-based serialization: how the Pickle virtual machine actually works, how an attacker can craft a .pt file that silently runs arbitrary code while returning the correct tensor, and what weights_only=True actually protects against and where its limits are. We then look at Safetensors, Hugging Face's alternative format, which sidesteps the problem entirely by replacing Pickle with a format so simple you can write your own parser in half an hour. Along the way we cover the real performance differences between the two approaches and what the right migration path looks like for existing codebases.
Session Outline
- Introduction (3 min) : PyTorch 2.6 broke a lot of code. If you hit UnpicklingError after upgrading and just added weights_only=False to make it go away, this talk is for you.
- How PyTorch Serialization Works (4 min) : A quick walkthrough of torch.save and torch.load using a minimal example, leading to the reveal that pickle is doing all the actual work under the hood and why that matters.
- Deep Dive: The Pickle Virtual Machine (6 min) : Pickle is not just a data format. It is a small programming language that runs on its own virtual stack machine. Using pickletools, we disassemble real pickle bytes and step through op codes like GLOBAL, MARK, UNICODE, TUPLE, REDUCE, and STOP to build a precise mental model of what deserialization actually does.
- Live Demo: Crafting a Malicious Tensor (5 min) : We build a .pt file that silently executes arbitrary code while returning the correct tensor back to the caller. No error, no warning, the right output. This section exists to make the threat concrete rather than theoretical, and to show exactly what weights_only=True is protecting against.
- Mitigation Strategies (5 min) : Three options and an honest assessment of each:
Security scanners (fickling, ProtectAI, ModelScan): useful as a first filter, but operating on a blacklist that motivated attackers can work around.
weights_only=True: now the PyTorch 2.6 default, which is genuinely good news. But the restricted unpickler is still a restricted pickle runtime, and restricted runtimes have been bypassed before in other Python contexts.
SafeTensors: drops pickle entirely and is the only option that removes the attack surface rather than just shrinking it. - SafeTensors Deep Dive (4 min) : The full format explained: an 8-byte little-endian length prefix, a UTF-8 JSON header mapping tensor names to dtype, shape, and byte offsets, followed by raw tensor data. We write a working parser from scratch in roughly 30 lines of Python to show just how auditable the format actually is compared to pickle.
- Summary & Q&A (3 min)
Prerequisites
Basic Python and some exposure to PyTorch or any tensor library. Minimal security experience needed. The attack concepts are built up from scratch.
Methodology
The talk is code-first throughout, using only the Python standard library and PyTorch:
- pickle and pickletools are used to serialise a simple Python tuple, disassemble the resulting bytes, and walk through each op code on a virtual stack. This is the foundation needed to understand both the attack and the mitigation.
- A malicious .pt file is assembled by hand: a GLOBAL(builtins.eval) + REDUCE payload is prepended to a legitimate tensor file. Loading this file returns the correct tensor while silently executing arbitrary code, demonstrating exactly the class of attack that motivated the PyTorch 2.6 default change.
- weights_only=True is tested against the same malicious file and raises UnpicklingError: unsupported global: builtins.eval was not allowed. We then dig into what the restricted unpickler actually allows, how the safe_globals allowlist and TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD environment variable work for legacy migration, and why "restricted pickle" is still pickle.
- The SafeTensors format is parsed manually using struct and json. To show it's simplicity
Key Takeaways
After this session, attendees will be able to:
- Explain the reason behind the PyTorch 2.6 breaking change and why adding weights_only=False to silence the error is the wrong response in most situations.
- Read disassembled pickle byte streams using pickletools and identify dangerous op codes like GLOBAL and REDUCE.
- Understand the full attack chain: how a malicious tensor file is constructed, why it executes silently, and why the victim sees a completely valid tensor with no indication anything went wrong.
- Assess the actual security boundary of weights_only=True: what it blocks, what the allowlist mechanism is for, and why it reduces rather than eliminates the attack surface.
- Parse the SafeTensors format from scratch in pure Python, and understand why its simplicity is a security property rather than a limitation.
Gregor Ziegltrum is a Senior Consultant at TNG Technology Consulting, specializing in innovative AI solutions. His work spans building agentic systems, as well as benchmarking, training, and fine-tuning Large Language Models. Beyond project work, Gregor shares his expertise through talks and internal trainings on Python best practices and deep dives into quirky language features. Before joining TNG, he earned both his Bachelor’s and Master’s degrees in Informatics from TU Munich, where he developed a strong passion for Data, AI, and High-Performance Computing.