Lightning-AI / Lightning-AI/pytorch-lightning
Resume training from checkpoint that only save trainable parameters
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
My goal is to reduce the file size
For example, saving a checkpoint with [SONAR model ](https://github.com/facebookresearch/SONAR) require 6GB of disk space
My solution to the problem above is to replace `state_dict` from checkpoint with only trainable dict
```
class LitModel(LightningModule)
def on_save_checkpoint(self, checkpoint):
checkpoint['state_dict'] = self.get_trainable_state_dict()
def get_trainable_state_dict(self):
state = {}
for name, param in self.named_parameters():
if param.requires_grad:
state[name] = param.data.cpu()
for name, buffer in self.named_buffers():
state[name] = buffer.data.cpu()
return state
```
However, one thing that I notice is that I require `strict=False` in `LitModel.load_from_checkpoint(strict=False)` to load the checkpoint. So I assume resuming from checkpoint using `trainer.fit(ckpt_path)` would also fail without `strict=False`
PS: I have not tried it, I am in the middle of training (3.8 out of 13 hours), I don't want to risk it
### Pitch
_No response_
### Alternatives
_No response_
### Additional context
_No response_
cc @lantiga @borda
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 tracing checkpoint loading from trainer.fit(ckpt_path) and compare it with LitModel.load_from_checkpoint(strict=False). Review how the on_save_checkpoint state_dict replacement interacts with missing parameters and resume state. Done means the expected behavior for checkpoints containing only trainable parameters is established and covered by an appropriate test.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100