Lightning-AI / Lightning-AI/pytorch-lightning

Feature request: Support splitting model weights and training states into separate checkpoint files

Open
#21,170 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

callback: model checkpoint checkpointing feature
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Description & Motivation

# 🚀 Feature Request: Support splitting model weights and training states into separate checkpoint files

## Feature Request

Currently, PyTorch Lightning saves the entire training state (model weights, optimizer states, scheduler states, trainer state, etc.) into a single `.ckpt` file.

I would like to have an option to **separate model weights (and config) from training states** when saving checkpoints.

For example, the desired checkpoint structure could look like this:

```
checkpoints/
pretrained_model/
config.json # model configuration
model.safetensors # model weights only
training_states.pth # optimizer, LR scheduler, trainer states
```

### Pitch

_No response_

### Alternatives

_No response_

### Additional context

_No response_

cc @lantiga @borda

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 locating PyTorch Lightning’s existing checkpoint save and load entry points, since no files or tests are named. Compare the current single .ckpt structure with the requested config.json, model.safetensors, and training_states.pth layout. Done should include a defined option for separate files and a clear way to restore training state.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.