Skip to content

fix: load checkpoints with torch.load(weights_only=True)#2287

Open
Solaris-star wants to merge 1 commit into
Lightning-AI:mainfrom
Solaris-star:fix/2189-weights-only-true
Open

fix: load checkpoints with torch.load(weights_only=True)#2287
Solaris-star wants to merge 1 commit into
Lightning-AI:mainfrom
Solaris-star:fix/2189-weights-only-true

Conversation

@Solaris-star

Copy link
Copy Markdown

Summary

Several production paths loaded .pth checkpoints with weights_only=False or without weights_only, which allows pickle gadget execution from untrusted files.

Set weights_only=True for model weight loads in:

  • litgpt/api.py
  • litgpt/utils.py (load_checkpoint ModelParallel path)
  • litgpt/scripts/convert_pretrained_checkpoint.py
  • litgpt/scripts/merge_lora.py
  • litgpt/generate/sequentially.py / tp.py
  • litgpt/eval/evaluate.py conversion helper

Testing

  • Static review of remaining production torch.load call sites
  • Full suite not run here (host lacks torch/litgpt test env); CI should cover

Fixes #2189

Untrusted .pth checkpoints can execute arbitrary code via pickle when
loaded with weights_only=False (or omitted). Prefer weights_only=True for
model weight loads in the API, utils, merge/convert scripts, generate
paths, and eval conversion helper.

Fixes Lightning-AI#2189

Signed-off-by: Solaris-star <820622658@qq.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Unsafe Checkpoint Loading - This Is Bad

1 participant