Lightning-AI / Lightning-AI/pytorch-lightning

Handle LazyModuleMixin Properly

Open
#13,764 10 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature pl
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🚀 Feature

Is it possible to add an optional initialize_loop prior to configure_optimizer and compute model summary? The initialize_loop shall be supplied with a batch of training data that can be used to run through `LazyModuleMixin` so that unintialized parameters can be initialized.

### Motivation

I don't know start from which version of PyTorch, it introduces a convenient way to construct a module that eliminates the need to calculate input shapes for each individual layer. However, the tradeoff is that it relies on passing one dummy batch of data to the model prior to configure optimizer and training the model.

Currently, the `Trainer.fit` does things in the following order:

1. attach data loader
2. configure accelerator strategy
3. call `model.on_fit_start()`
4. call `model.configure_optimizers()`
5. if sanity_check, run validation loop with `num_sanity_val_steps` batches of validation data.
6. start running training loop.
7. call `model.on_fit_end()`

However, the current flow throws `UninitializedParameter` error when configure_optimizers method is called. To get around the issue, currently I will have to initialize the parameters outside `pl.Trainer` by loading one batch of training data module to the model on my own, which involves manually setting up which accelerator I should for the initialization loop.

### Pitch

I suggest adding an additional hook for the initialization loop and let lightning takes care of setting up proper accelerator and supplying one batch of training data. The proposed new workflow for `Trainer.fit` is presented as follows:

1. attach data loader
2. configure accelerator strategy
3. call `model.on_fit_start()`
4. **run one optional initialization loop to initialize module parameters.**
5. call `model.configure_optimizers()`
6. if sanity_check, run validation loop with `num_sanity_val_steps` batches of validation data.
7. start running training loop.
8. call `model.on_fit_end()`

### Alternatives

Without the new features, I will have to write initialization loop that looks like the following:

```python
model = ModelWithLazyModule()
for batch in train_loader:
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
batch = batch.to(device)
model(batch)
break
trainer = pl.Trainer()
trainer.fit(model, train_loader)
```

### Additional context

To my knowledge, people mentioned `UninitializedParameter` in #7546 and #7642. It seems to focus on computing model summary table if I understand correctly.

Also, Lightning backend is complicated to me and I know I am not capable of implementing this feature on my own.

cc @borda @akihironitta

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 the Trainer.fit lifecycle described in the issue, especially on_fit_start(), configure_optimizers(), the training loader, and the accelerator strategy. Determine how an optional initialization loop could consume one training batch before optimizer configuration and model summary; done means LazyModuleMixin parameters are initialized without manual device setup while existing training and sanity-check behavior remains intact.

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
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.