| # 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: |
| ```python |
| 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)) |
| ``` |
|
|
