Flash Attention Training Framework Insecure Deserialization Vulnerability Allowing Arbitrary Code Execution

A vulnerability allowing insecure deserialization has been identified in the Flash Attention training framework, specifically in versions prior to the commit on April 13, 2025. This vulnerability arises in the checkpoint loading mechanism, where the load_checkpoint() function in checkpoint.py and the checkpoint loading code in eval.py use torch.load() without the security-restrictive weights_only=True parameter. As a result, arbitrary Python objects can be deserialized via the pickle module. An attacker could exploit this by providing a maliciously crafted checkpoint file, which, when loaded during model warmstarting or evaluation, would execute arbitrary code on the victim's system.

Impact

Exploitation of this vulnerability allows for arbitrary code execution on the system of the user loading the compromised checkpoint file.

Reproduction

To reproduce this vulnerability, load a maliciously crafted checkpoint file into the Flash Attention training framework using the load_checkpoint() function. The checkpoint file must be designed to exploit the insecure deserialization by including arbitrary Python objects that, when deserialized, execute code on the system.