WandbLogger in DDP setting creates multiple runs in W&B

[konstantinjdobler]

When using the WandbLogger in DDP mode with n>1 GPUs, n runs in Weights&Biases will be created. Only the first run will have the actual logs, all other runs will be empty, essentially cluttering up the W&B project.

Expected behavior

Only one run will be created, even when using DDP on multiple GPUs.

Environment

Example config, but this issue also happens on other configurations with different GPU models:

• CUDA:
  • GPU:
    • NVIDIA A100-SXM-80GB
    • NVIDIA A100-SXM-80GB
  • available: True
  • version: 11.0
• Packages:
  • numpy: 1.22.0
  • pyTorch_debug: False
  • pyTorch_version: 1.10.0
  • pytorch-lightning: 1.5.8
  • tqdm: 4.62.3
• System:
  • OS: Linux
  • architecture:
    • 64bit
    • ELF
  • processor: x86_64
  • python: 3.9.9
  • version: # 91-Ubuntu SMP Thu Jul 15 19:09:17 UTC 2021

Additional Info

Since this bug only happens in DDP with multiple GPUs, I can't reproduce it in Colab.

cc @awaelchli @morganmcg1 @AyushExel @borisdayma @scottire

[konstantinjdobler]

I think I can "fix" this by manually calling wand.init only on the rank0 process once before creating the WandbLogger, but it would be much more in the spirit of pytorch-lightning if the user didn't have to think about these things.

EDIT: Some more information:
My WandbLogger looks like this and I am calling wandb.watch afterwards:

wandb_logger = WandbLogger(
        project=<project>, entity=<entity>, name=<name>, tags=<tags>, log_model=True)
wandb_logger.watch(model, log="all", log_graph=False)

I do not make any manual calls to wandb in my project, I only use the integration exposed through pytorch-lightning.

[morganmcg1]

Hey @konstantinjdobler, Morgan from W&B here, thanks for the issue. Yep like you found, calling wand.init on the rank0 process is one way to work with DDP today.

We also have an experimental feature called wandb service that improves wandb's handling of multiprocessing. You can read more about here. It might address your issue if you'd like to try?

Note that as part of it you'll have to install this branch:

pip install --force git+https://github.com/wandb/pytorch-lightning.git@wandb-service-attach

[konstantinjdobler]

I tried wandb service but unfortunately the issue is not resolved by it.

After a bit of investigation, I believe the issue runs much deeper and only happens when running jobs via SLURM.
In the WandbLogger, a new wandb run will be created in the experiment method. The @rank_zero_experiment decorator states: "Returns the real experiment on rank 0 and otherwise the DummyExperiment."
https://github.com/PyTorchLightning/pytorch-lightning/blob/cec4f17df3cbdf2b585b723d518e03bd500a7c5b/pytorch_lightning/loggers/wandb.py#L320-L322

But here is the issue! There seems to be a bug in the calculation for rank 0. The _get_rank method searches the os.environ for the keys "RANK", "SLURM_PROCID", "LOCAL_RANK" and stops at the first hit. When I submit my slurm jobs, SLURM_PROCID is always set (to 0) and therefore supersedes the correct value in LOCAL_RANK.

https://github.com/PyTorchLightning/pytorch-lightning/blob/cec4f17df3cbdf2b585b723d518e03bd500a7c5b/pytorch_lightning/utilities/distributed.py#L46-L67

When I simply swap the order of both keys, the issue is resolved (probably because the @rank_zero_experiment decorator is working correctly again). Now SLURM_PROCID could have other values than 0, but given my example and general reasoning, I don't see why LOCAL_RANK shouldn't supersede SLURM_PROCID.

In fact, I remember not having this issue using pytorch-lightning==1.2.0 and in 1.2.8 in this commit 3fb7c5e the change was made. I created a PR #11406 to fix this.

[function2-llx]

There's still some issue with this. _get_rank will check LOCAL_RANK, while wandb seems to assume rank_zero_only.rank to be the global rank, e.g., as the following comment suggests:
https://github.com/PyTorchLightning/pytorch-lightning/blob/f9b69ce5b0de32bfd2300eb8028fdbebf23fe5a1/pytorch_lightning/loggers/wandb.py#L378

Then it will still create multiple runs on multiple nodes, as it will find exact one process with LOCAL_RANK of zero on each node.

[awaelchli]

@function2-llx In which cluster environment are you? How are you launching the processes?

[function2-llx]

@awaelchli I treat my cluster as a general-purpose cluster, and I launch the processes by following the doc of general-purpose cluster, i.e., firstly configuring environment variables on each node, and executing the same script on each node.

Currently, I made some custom fix to make the logger create a run only for process with NODE_RANK=0 and LOCAL_RANK=0.

def node_zero_only(fn: Callable) -> Callable:
    @wraps(fn)
    def wrapped_fn(*args, **kwargs) -> Optional[Any]:
        if node_zero_only.node == 0:
            return fn(*args, **kwargs)
        return None
    return wrapped_fn

node_zero_only.node = getattr(node_zero_only, 'node', int(os.environ.get('NODE_RANK', 0)))

def node_zero_experiment(fn: Callable) -> Callable:
    """Returns the real experiment on rank 0 and otherwise the DummyExperiment."""
    @wraps(fn)
    def experiment(self):
        @node_zero_only
        def get_experiment():
            return fn(self)
        return get_experiment() or DummyExperiment()
    return experiment

class MyWandbLogger(WandbLogger):
    @WandbLogger.experiment.getter
    @node_zero_experiment
    def experiment(self):
        return super().experiment

[morganmcg1]

Hi, @function2-llx , @awaelchli , getting back to this, are you running Lightning v1.6? Wandb service is turned on by default there. I know you mentioned its not solved by service so I'll raise this with out eng team

[function2-llx]

Hello @morganmcg1 , yes I've tested with PyTorch Lightning 1.6.0 (2 weeks ago).

[undertherain]

I have the same problem with PL 1.6.3
I use custom cluster environment, launching an MPI process for each GPU
I set environment RANK variable manually in each process before creating WANDB logger but it does not seem to help.
consider reopen?

[undertherain]

btw in rank_zero.py:
rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK")
how does checking for LOCAL_RANK make sense? Shouldn't it be global rank?

[undertherain]

okey... and it didn't work because rank zero is checked on import
I got around it it by
pl.utilities.rank_zero.rank_zero_only.rank = cluster_env.local_rank()
before doing anything with PL but this is IMO pretty ugly and I really hope devs would consider reopening this issues and fixing it properly.