Lightning-AI / Lightning-AI/pytorch-lightning
`Trainer.validate()` after `Trainer.fit()` not working with FSDP and `auto_wrap_policy`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
I'm training large model with FSDP. Moreover, I'm using `_HYBRID_SHARD_ZERO2` sharding strategy and specify `configure_model` methon in Lightning Module.
At the moment I explicitly call validation after fit:
```python
trainer.fit(my_model, train_dataloader, val_dataloaders)
trainer.validate(my_model, val_dataloaders)
```
The fitting phase is fine, but when calling the validation I encounter this warning:
```python
A FSDP `auto_wrap_policy` is set, but the model is already wrapped. The policy will be ignored.
```
And after that the validation fails due to an error in `torch.Embedding` (RuntimeError: "weight" must be 2-D).
My guess is that after ignoring the policy, the model is sharded incorrectly and some nodes are running with the wrong embedding weights 🤷♂️ This is probably related to #18971, but I get the same error on the last version (2.1.2)
### What version are you seeing the problem on?
v2.1
### How to reproduce the bug
_No response_
### Error messages and logs
```
# Error messages and logs here please
```
### Environment
Current environment
```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```
### More info
_No response_
cc @awaelchli @carmocca
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 with the reported Trainer.fit() followed by Trainer.validate() sequence, using FSDP, _HYBRID_SHARD_ZERO2, auto_wrap_policy, and a LightningModule with configure_model. The issue provides no reproduction files or tests, so first isolate the failure and the model-wrapping state; done means validation completes without the auto_wrap_policy warning or the torch.Embedding shape error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100