WandbLogger: log distributed training experiments (multi-node)

[guyi2000]

Description & Motivation

Currently, when using PyTorch Lightning with WandbLogger for multi-node distributed training, only the system metrics (CPU/GPU utilization, memory usage, etc.) from node 0 and logs from rank 0 are recorded by wandb. This limits visibility into critical performance data from non-zero ranks, making it harder to debug issues like uneven resource utilization across nodes/GPUs or pinpoint hardware bottlenecks in distributed setups.

With wandb's recently added distributed experiment tracking support (also the end-to-end example report), it's now possible to collect system metrics and logs from all nodes. However, Lightning's WandbLogger does not yet leverage this capability out-of-the-box. This feature request proposes updating Lightning's wandb integration to support full distributed system monitoring.

Pitch

Add native support in pytorch_lightning.loggers.WandbLogger to:

• Log system metrics (hardware telemetry) from ​​all nodes​​ in distributed training, not just node 0

This could be implemented by:

• Adding a log_all_ranks: bool parameter to WandbLogger to enable/disable this behavior
• wandb.init(..., settings=wandb.Settings(x_label="rank_0", mode="shared", x_primary=True)) for rank 0
• wandb.init(..., settings=wandb.Settings(x_label=f"rank_{rank}", mode="shared", x_primary=False)) for non-zero ranks

Alternatives

While users could manually override logging behavior by modifying wandb.init() parameters in non-zero ranks, this:

• Conflicts with Lightning's logger orchestration
• Risks creating multiple wandb runs unintentionally
• Requires error-prone custom code outside Lightning's abstractions

A cleaner native implementation would provide better safety and usability.

Additional context

wandb has recently updated the feature, here is the relevant issue:

• [App]: hardware logging in multi-node setting wandb/wandb#7470 (comment)

When examining the WandbLogger source code, the @rank_zero_experiment decorator enforces wandb initialization and logging exclusively on rank 0. This creates architectural challenges for multi-rank logging because:

• Non-zero ranks never initialize a wandb.Run
• Any attempt to modify the existing logger would require significant refactoring of the decorator logic

Instead of adding a log_all_ranks parameter to WandbLogger, a cleaner solution could be:

• Create a new WandbDistributedLogger class under lightning.pytorch.loggers.wandb
• Design this class to:
  • Initialize wandb runs on ​​all ranks​​ using wandb.init(..., settings=wandb.Settings(x_label=..., mode="shared", x_primary=...))
  • Bypass the @rank_zero_experiment restriction
  • Preserve the original WandbLogger behavior (e.g. log_metrics, log_hyperparams, etc.)

This approach:

• Maintains backward compatibility
• Avoids complicating the existing WandbLogger API

I would be happy to submit a PR implementing if needed.

• The new WandbDistributedLogger class
• Integration tests verifying multi-node metric collection

cc @lantiga @Borda @morganmcg1 @borisdayma @scottire @parambharat

[guyi2000]

I've prototyped a WandbDistributedLogger implementation (full code) that demonstrates the core functionality:

import os
import wandb
import psutil

from wandb.sdk.lib import RunDisabled
from wandb.wandb_run import Run

from lightning_utilities.core.imports import RequirementCache
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch.utilities.rank_zero import rank_zero_warn

from typing import Union, Optional

import torch.distributed as dist


class WandbDistributedLogger(WandbLogger):
    def __init__(self, *args, **kwargs) -> None:
        _WANDB_AVAILABLE = RequirementCache("wandb>=0.19.9")
        if not _WANDB_AVAILABLE:
            raise ModuleNotFoundError(str(_WANDB_AVAILABLE))

        super().__init__(*args, **kwargs)

        self._experiment: Optional[Union[Run, RunDisabled]]  # Add this for type-checking tools (mypy)

    @property
    def experiment(self) -> Union["Run", "RunDisabled"]:
        r"""Actual wandb object. To use wandb features in your :class:`~lightning.pytorch.core.LightningModule` do the
        following.

        Example::

        .. code-block:: python

            self.logger.experiment.some_wandb_function()

        """
        if not (dist.is_available() and dist.is_initialized()):
            raise RuntimeError(
                "`WandbDistributedLogger` requires DDPStrategy training to be initialized."
            )

        if self._experiment is None:
            local_rank = dist.get_node_local_rank()
            global_rank = dist.get_rank()
            world_size = dist.get_world_size()

            if self._offline:
                os.environ["WANDB_MODE"] = "dryrun"

            attach_id = getattr(self, "_attach_id", None)
            if wandb.run is not None:
                # wandb process already created in this instance
                rank_zero_warn(
                    "There is a wandb run already in progress and newly created instances of `WandbLogger` will reuse"
                    " this run. If this is not desired, call `wandb.finish()` before instantiating `WandbLogger`."
                )
                self._experiment = wandb.run
            elif attach_id is not None and hasattr(wandb, "_attach"):
                # attach to wandb process referenced
                self._experiment = wandb._attach(attach_id)
            else:
                # create new wandb process
                if global_rank == 0:
                    settings = self._wandb_init.pop("settings", None)
                    if settings is None:
                        settings = wandb.Settings()
                    settings.update_from_dict(
                        {
                            "mode": "shared",
                            "x_label": "rank-0",
                            "x_primary": True,
                            "x_update_finish_state": True,
                            "x_stats_gpu_device_ids": [local_rank],
                        }
                    )
                    self._experiment = wandb.init(**self._wandb_init, settings=settings)
                    run_id = [self._experiment.id]
                else:
                    run_id = [""]
                # non-zero ranks need rank 0 run id
                dist.broadcast_object_list(run_id)
                if global_rank != 0:
                    self._wandb_init["id"] = run_id[0]
                    settings = self._wandb_init.pop("settings", None)
                    if settings is None:
                        settings = wandb.Settings()
                    settings.update_from_dict(
                        {
                            "mode": "shared",
                            "x_label": f"rank-{global_rank}",
                            "x_primary": False,
                            "x_update_finish_state": False,
                            "x_stats_gpu_device_ids": [local_rank],
                        }
                    )
                    self._experiment = wandb.init(**self._wandb_init, settings=settings)

                if isinstance(self._experiment, (Run, RunDisabled)) and self._experiment._metadata is not None:
                    self._experiment._metadata.gpu_count = world_size
                    self._experiment._metadata.cpu_count = psutil.cpu_count(
                        False
                    ) * int(os.environ.get("NODE_COUNT", 1))
                    # We cannot get num nodes in dist, but many cluster will set some environment variables
                    self._experiment._metadata.cpu_count_logical = psutil.cpu_count(
                        True
                    ) * int(os.environ.get("NODE_COUNT", 1))

                # define default x-axis
                if isinstance(self._experiment, (Run, RunDisabled)) and getattr(
                    self._experiment, "define_metric", None
                ):
                    if self._wandb_init.get("sync_tensorboard"):
                        self._experiment.define_metric("*", step_metric="global_step")
                    else:
                        self._experiment.define_metric("trainer/global_step")
                        self._experiment.define_metric(
                            "*", step_metric="trainer/global_step", step_sync=True
                        )

        return self._experiment

Current Implementation Constraints​​:

• Only suitable for DDPStrategy (only tested with NCCL backend)
• Initializes wandb runs on ​​every rank processes per node
• CPU metrics collection using psutil (needs better hardware abstraction)
• Requires wandb>=0.19.9 for shared run tracking

This serves as a starting point, and the code above is not robust enough - I welcome maintainer feedback on architectural decisions and would be happy to refine this implementation through a PR.