Lightning-AI / Lightning-AI/pytorch-lightning
Parameters and Gradient is not logged by WandB under FSDP strategy
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 find that when using FSDP strategy, the model parameters and gradients are not logged by WandB. However, everything works well if I switch FSDP to native DDP strategy.
Since the gradients are hooked by wandb.run.watch, I am not sure this is a lighting issue or wandb issue.
### What version are you seeing the problem on?
master
### 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: 2.1.0dev (updated to bleeding-edge)
#- WandB Version: 0.14.0
```
### More info
_No response_
cc @lantiga @morganmcg1 @borisdayma @scottire @parambharat
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 by reproducing the report with PyTorch Lightning 2.1.0dev, WandB 0.14.0, and the FSDP strategy, then compare it with native DDP. Trace the WandB watch integration and gradient or parameter logging behavior; done means both are logged under FSDP as they are under DDP.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100