Lightning-AI / Lightning-AI/litgpt

Unsafe Checkpoint Loading - This Is Bad

Open
#2,189 1 comment 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.