Lightning-AI / Lightning-AI/pytorch-lightning

Resume training from checkpoint that only save trainable parameters

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

Nobody has claimed this yet.

checkpointing feature
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.