Lightning-AI / Lightning-AI/litgpt
Unsafe Checkpoint Loading - This Is Bad
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 1.5k
- Avg merge
- 15h 37m
- Merged PRs (30d)
- 1
Description
**Severity:** Critical Security Issue
Alright, so this one's pretty serious. The codebase loads model checkpoints with `weights_only=False` all over the place:
```python
# litgpt/api.py, line 117
state_dict = torch.load(self.checkpoint_dir / "lit_model.pth", weights_only=False)
```
### Why this matters:
Pickle files (what PyTorch uses) can execute arbitrary code. Like, literally anything. Someone shares a "fine-tuned model" on HuggingFace, you download it, boom - they're mining crypto on your GPU cluster. Or worse.
PyTorch added the `weights_only` flag specifically because of this. There's a reason they made it - use it.
### What I found:
- `litgpt/api.py` lines 117, 397, 421 - all unsafe
- `litgpt/utils.py` line 393 - uses mmap but no weights_only check
- Multiple converter scripts do the same thing
- Even the tutorial code shows unsafe loading patterns
### What needs to happen:
1. Change ALL `torch.load()` calls to use `weights_only=True`
2. The codebase already depends on safetensors - just use that as the primary format
3. Add a verification step for downloaded checkpoints (checksums at minimum)
4. Put a big fat warning in the docs about loading untrusted checkpoints
This isn't theoretical. People WILL download random checkpoints from the internet. Make it safe by default.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by inventorying the torch.load calls in litgpt/api.py, litgpt/utils.py, converter scripts, and tutorial code, then review the existing checkpoint and safetensors handling. Done means unsafe loading is removed or explicitly protected across the codebase, downloaded checkpoints have verification, and documentation warns about untrusted checkpoints.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- ai, security
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100