accelerate-ace-poc / README.md
etwithin's picture
Add ACE PoC
bcd4e63 verified

accelerate load_custom_state() ACE PoC

This checkpoint demonstrates arbitrary code execution via accelerate's load_custom_state() function which explicitly uses weights_only=False.

When loaded via accelerator.load_state(), the pickle payload executes arbitrary code before any validation.

Vulnerable code in accelerate/checkpointing.py:

def load_custom_state(obj, path, index: int = 0):
    load_location = f"{path}/custom_checkpoint_{index}.pkl"
    obj.load_state_dict(load(load_location, map_location="cpu", weights_only=False))