lendle
commented
1 year ago I reported this to wandb a couple of months ago. Here's a workaround I've been using, which I think should also work when not using FSDP:
``
In your LightningModule, implement
def on_before_optimizer_step(self, optimizer):
if self.trainer._logger_connector.should_update_logs:
with FullyShardedDataParallel.summon_full_params(
self.get_configured_model(),
rank0_only=True,
writeback=False,
with_grads=True,
):
if self.global_rank == 0:
if isinstance(self.logger, WandbLogger):
for name, param in self.model.named_parameters():
torch_history: TorchHistory = self.logger.experiment._torch
torch_history.log_tensor_stats(param, f"parameters/{name}")
if param.requires_grad:
torch_history.log_tensor_stats(
param.grad, f"gradients/{name}"
)
stale[bot]
konstantinjdobler
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
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