Lightning-AI / Lightning-AI/pytorch-lightning
Handle LazyModuleMixin Properly
Nobody has claimed this yet.
- 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
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 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