Unsafe Checkpoint Loading - This Is Bad

[AI-God-Dev]

Severity: Critical Security Issue

Alright, so this one's pretty serious. The codebase loads model checkpoints with weights_only=False all over the place:

# 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.