CVE-2026-31221: pytorch-lightning: RCE via insecure checkpoint deserialization

AWAITING NVD
Published May 12, 2026
CISO Take

PyTorch-Lightning versions ≤2.6.0 allow arbitrary code execution when loading model checkpoints because load_from_checkpoint() calls torch.load() without the weights_only=True parameter, enabling Pickle to deserialize arbitrary Python objects with no restrictions. Any ML pipeline that loads checkpoints from an attacker-influenced source — shared model hubs, third-party artifact stores, collaborative research environments, or CI/CD pipelines pulling external weights — is at risk of full system compromise requiring no authentication beyond getting the victim to open the file. No CVSS score or EPSS data is available yet and no public exploit exists, but CWE-502 in ML checkpoint loading is trivially exploitable by anyone with the ability to deliver a crafted .ckpt file, and pytorch-lightning is ubiquitous in production training and fine-tuning workflows. Until an official patch is released, restrict checkpoint loading to cryptographically verified, trusted sources and evaluate migrating weight storage to the safetensors format which eliminates Pickle deserialization risk entirely.

Sources: NVD ATLAS GitHub Advisory

Risk Assessment

HIGH. Insecure Pickle deserialization in ML checkpoints is a well-understood, trivially exploitable attack class — crafting a malicious checkpoint requires only basic Python knowledge and no AI/ML expertise. The vulnerability affects any system that loads pytorch-lightning checkpoints from an untrusted or attacker-influenced source, which is routine in ML workflows involving model sharing, fine-tuning from public checkpoints, or collaborative research. The absence of CVSS scoring reflects an NVD data gap, not low severity. PyTorch-Lightning is extensively deployed in production ML training and inference pipelines globally, and GPU training nodes frequently run with elevated privileges and access to valuable assets including API keys, proprietary datasets, and model weights.

Attack Kill Chain

Artifact Preparation
Adversary crafts a malicious PyTorch-Lightning checkpoint (.ckpt) file with an embedded Pickle payload that executes arbitrary Python code upon deserialization.
AML.T0018.002
Delivery
Adversary distributes the malicious checkpoint via a public model hub, compromised artifact repository, or social engineering such as sharing a purported 'fine-tuned model' with a target team.
AML.T0011.000
Exploitation
Victim calls LightningModule.load_from_checkpoint() on the malicious file; torch.load() deserializes the Pickle payload without restrictions, executing attacker-controlled code on the ML host.
AML.T0112.001
Impact
Attacker achieves RCE on ML infrastructure, enabling exfiltration of model weights, API keys, and training data; persistent backdoor installation; or lateral movement to adjacent production systems.
AML.T0025

Affected Systems

Package Ecosystem Vulnerable Range Patched
pytorch-lightning pip No patch
31.1K OpenSSF 5.2 1.6K dependents Pushed 3d ago 40% patched ~496d to patch Full package profile →

Do you use pytorch-lightning? You're affected.

Severity & Risk

CVSS 3.1
N/A
EPSS
N/A
Exploitation Status
No known exploitation
Sophistication
Trivial

Recommended Action

1 step
  1. 1) Upgrade PyTorch-Lightning to ≥2.6.1 when a patch is released; monitor the upstream advisory at github.com/Lightning-AI/pytorch-lightning for release timing. 2) Until patched, audit all load_from_checkpoint() call sites and enforce that checkpoint files are loaded exclusively from trusted, integrity-verified sources using SHA-256 checksums or signed artifact manifests. 3) As a code-level workaround, replace internal torch.load() calls in custom checkpoint loading logic with torch.load(f, weights_only=True) where direct torch access is available. 4) Migrate model weight storage to safetensors format, which eliminates Pickle deserialization risk by design. 5) Apply network egress controls on ML training infrastructure to detect reverse shell callbacks as an indicator of successful exploitation. 6) Scan artifact repositories for unexpected checkpoint files that may have been substituted by an attacker targeting your model supply chain.

Classification

Compliance Impact

This CVE is relevant to:

EU AI Act
Article 15 - Accuracy, robustness and cybersecurity
ISO 42001
A.6.1.5 - AI system supply chain
NIST AI RMF
GOVERN 1.1 - Organizational AI risk policies and procedures MANAGE 2.4 - AI risk treatment and response mechanisms
OWASP LLM Top 10
LLM03:2025 - Supply Chain Vulnerabilities

Frequently Asked Questions

What is CVE-2026-31221?

PyTorch-Lightning versions ≤2.6.0 allow arbitrary code execution when loading model checkpoints because load_from_checkpoint() calls torch.load() without the weights_only=True parameter, enabling Pickle to deserialize arbitrary Python objects with no restrictions. Any ML pipeline that loads checkpoints from an attacker-influenced source — shared model hubs, third-party artifact stores, collaborative research environments, or CI/CD pipelines pulling external weights — is at risk of full system compromise requiring no authentication beyond getting the victim to open the file. No CVSS score or EPSS data is available yet and no public exploit exists, but CWE-502 in ML checkpoint loading is trivially exploitable by anyone with the ability to deliver a crafted .ckpt file, and pytorch-lightning is ubiquitous in production training and fine-tuning workflows. Until an official patch is released, restrict checkpoint loading to cryptographically verified, trusted sources and evaluate migrating weight storage to the safetensors format which eliminates Pickle deserialization risk entirely.

Is CVE-2026-31221 actively exploited?

No confirmed active exploitation of CVE-2026-31221 has been reported, but organizations should still patch proactively.

How to fix CVE-2026-31221?

1) Upgrade PyTorch-Lightning to ≥2.6.1 when a patch is released; monitor the upstream advisory at github.com/Lightning-AI/pytorch-lightning for release timing. 2) Until patched, audit all load_from_checkpoint() call sites and enforce that checkpoint files are loaded exclusively from trusted, integrity-verified sources using SHA-256 checksums or signed artifact manifests. 3) As a code-level workaround, replace internal torch.load() calls in custom checkpoint loading logic with torch.load(f, weights_only=True) where direct torch access is available. 4) Migrate model weight storage to safetensors format, which eliminates Pickle deserialization risk by design. 5) Apply network egress controls on ML training infrastructure to detect reverse shell callbacks as an indicator of successful exploitation. 6) Scan artifact repositories for unexpected checkpoint files that may have been substituted by an attacker targeting your model supply chain.

What systems are affected by CVE-2026-31221?

This vulnerability affects the following AI/ML architecture patterns: Training pipelines, Model serving, MLOps/experiment tracking, Transfer learning workflows, Collaborative ML environments.

What is the CVSS score for CVE-2026-31221?

No CVSS score has been assigned yet.

Technical Details

NVD Description

PyTorch-Lightning versions 2.6.0 and earlier contain an insecure deserialization vulnerability (CWE-502) in the checkpoint loading mechanism. The LightningModule.load_from_checkpoint() method, which is commonly used to load saved model states, internally calls torch.load() without setting the security-restrictive weights_only=True parameter. This default behavior allows the deserialization of arbitrary Python objects via the Pickle module. A remote attacker can exploit this by providing a maliciously crafted checkpoint file, leading to arbitrary code execution on the victim's system when the file is loaded.

Exploitation Scenario

An adversary uploads a malicious .ckpt file to a public model hub such as Hugging Face or a compromised internal artifact registry, crafted with a Pickle payload that spawns a reverse shell to attacker-controlled infrastructure. An ML engineer or automated CI/CD pipeline runs LightningModule.load_from_checkpoint('path/to/malicious.ckpt') as part of a fine-tuning or model evaluation workflow — a completely routine operation. PyTorch-Lightning calls torch.load() without weights_only=True, the Pickle deserializer executes the embedded payload, and the attacker receives a shell on the training node. From there, the attacker exfiltrates proprietary model weights, API keys from environment variables, training datasets, or moves laterally to adjacent systems accessible from the ML compute environment.

Timeline

Published
May 12, 2026
Last Modified
May 12, 2026
First Seen
May 12, 2026

Related Vulnerabilities