Improve multiprocessing compatibility with WandbLogger

[awaelchli]

🚀 Feature

Weights and Biases is updating their logging backend to better support logging in multiple processes. This issue proposes a change to the WandbLogger implementation for an optimal integration.

Motivation

After #10896, the logger experiment consistently gets created at the first call to logger.experiment directly in the spawned worker processes. This is generally required because run objects for various third-party loggers are not pickleable. Secondly, depending on the logger backend, starting an experiment in the main process and logging in the worker processes could potentially create unwanted duplicate experiments/runs. However, Wandb has a special worker process running specifically just for logging. This service process can communicate with the other worker processes, but only if it gets created in the main process (this is a new feature being developed), which is no longer the case after #10896. Therefore we seek a solution to initiate that experiment on the main process instead, while still avoiding the pickle issue.

Pitch

1. Call wandb.init in the constructor of WandbLogger.
2. Wrap it with the rank_zero_only and ensure rank > 0 returns a mocked experiment
3. Ensure that the logger is still picklable

Alternatives

A alternative we discussed was to introduce a hook for the logger that runs before fit spawns processes. Here, WandbLogger could get a chance to instantiate their experiment.

Additional context

This will introduce a slight inconsistency with other logger implementations. However, it is always up to the specific logger to implement how they handle their experiment creation. As long as it adheres to the user facing api and experiments get created at the right time before users log anything, this change is imo acceptable.

cc @kptkin @raubitsj

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @Borda @awaelchli @morganmcg1 @AyushExel @borisdayma @scottire @manangoel99