Lightning-AI / Lightning-AI/pytorch-lightning

Support `assign` argument in `LightningModule.load_from_checkpoint` method

Open
#19,663 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

Currently, `LightningModule.load_from_checkpoint` only allows to pass `strict` argument to `torch.nn.Module.load_state_dict` method. There is also a new `assign` argument (https://github.com/pytorch/pytorch/pull/102212) that can be used to efficiently load model parameters from a checkpoint (see discussion in https://dev-discuss.pytorch.org/t/state-of-model-creation-initialization-seralization-in-pytorch-core/1240). This is a FR to allow passing of `assign` argument to `load_state_dict` similar to the `strict` argument.

If this is something desired I can open a PR.

### Pitch

Implementation can be similar to the `strict` argument. One consideration is to make it backward-compatible (any suggestions on how to?).

### 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 at LightningModule.load_from_checkpoint and trace where it calls torch.nn.Module.load_state_dict. Follow the existing strict argument handling and determine how assign can be passed while preserving compatibility with older PyTorch versions. Done means callers can provide assign for checkpoint loading without breaking existing strict-only usage.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.