Log distributed training experiments

Use W&B to log distributed training experiments with multiple GPUs.

During a distributed training experiment, you train a model using multiple machines or clients in parallel. W&B can help you track distributed training experiments so that you can monitor metrics, system utilization, and model behavior across all participating processes. Use this guide if you run distributed training jobs and want to log experiment data with W&B.

Based on your use case, track distributed training experiments using one of the following approaches:

• Track a single process: Track a rank 0 process (also known as a “leader” or “coordinator”) with W&B. This is a common solution for logging distributed training experiments with the PyTorch Distributed Data Parallel (DDP) class.
• Track multiple processes: For multiple processes, you can either:
  • Track each process separately using one run per process. You can optionally group them together in the W&B App UI.
  • Track all processes to a single run.

Concurrent connections

Each concurrent connection takes compute, memory, and network resources. Even empty client connections that don’t log metrics regularly push system metric updates, leading to slower performance when loading charts.

W&B recommends that you limit the maximum number of concurrent client connections as appropriate for your workload and that you monitor resource usage over time. W&B has tested with a hard limit of 300 concurrent client connections in Dedicated Cloud.

In Multi-tenant Cloud organizations, client connections for distributed training are subject to the same rate limits as regular training runs. On Teams and Enterprise plans, you receive higher rate limits than on the Free plan.

Track a single process

The following sections describe how to track values and metrics available to your rank 0 process. Use this approach to track only metrics that are available from a single process. Typical metrics include GPU/CPU utilization, behavior on a shared validation set, gradients and parameters, and loss values on representative data examples.

Within the rank 0 process, initialize a W&B run with wandb.init() and log experiments (wandb.Run.log()) to that run.

The following sample Python script (log-ddp.py) demonstrates one way to track metrics on two GPUs on a single machine using PyTorch DDP. PyTorch DDP (DistributedDataParallel in torch.nn) is a library for distributed training. The basic principles apply to any distributed training setup, but the implementation may differ.

The Python script does the following:

1. Starts multiple processes with torch.distributed.launch.
2. Checks the rank with the --local_rank command-line argument.
3. If the rank is set to 0, sets up wandb logging conditionally in the train() function.

if __name__ == "__main__":
    # Get args
    args = parse_args()

    if args.local_rank == 0:  # only on main process
        # Initialize wandb run
        run = wandb.init(
            entity=args.entity,
            project=args.project,
        )
        # Train model with DDP
        train(args, run)
    else:
        train(args)

Explore an example dashboard showing metrics tracked from a single process.

The dashboard displays system metrics for both GPUs, such as temperature and utilization.

However, the loss values as a function epoch and batch size were only logged from a single GPU.

Track multiple processes

Track multiple processes with W&B when you need metrics, logs, or artifacts from more than only the rank 0 process. The following sections describe two approaches:

• Track each process separately by creating a run for each process.
• Track all processes to a single run.

Track each process separately

The following sections describe how to track each process separately by creating a run for each process. Within each run you log metrics, artifacts, and so forth to their respective run. Call wandb.Run.finish() at the end of training, to mark that the run has completed so that all processes exit properly.

You might find it difficult to keep track of runs across multiple experiments. To mitigate this, provide a value to the group parameter when you initialize W&B (wandb.init(group='group-name')) to keep track of which run belongs to a given experiment. For more information about how to keep track of training and evaluation runs in experiments, see Group Runs.

Use this approach if you want to track metrics from individual processes. Typical examples include the data and predictions on each node (for debugging data distribution) and metrics on individual batches outside of the main node. This approach isn’t necessary to get system metrics from all nodes nor to get summary statistics available on the main node.

The following Python code snippet demonstrates how to set the group parameter when you initialize W&B:

if __name__ == "__main__":
    # Get args
    args = parse_args()
    # Initialize run
    run = wandb.init(
        entity=args.entity,
        project=args.project,
        group="DDP",  # all runs for the experiment in one group
    )
    # Train model with DDP
    train(args, run)

    run.finish()  # mark the run as finished

Explore the W&B App UI to view an example dashboard of metrics tracked from multiple processes. The two runs are grouped together in the left sidebar. Click a group to view the dedicated group page for the experiment. The dedicated group page displays metrics from each process separately.

The previous image demonstrates the W&B App UI dashboard. The sidebar shows two experiments: one labeled “null” and a second (bound by a yellow box) called “DPP”. Expand the group (select the Group dropdown) to see the runs associated with that experiment.

Organize distributed runs

Set the job_type parameter when you initialize W&B (wandb.init(job_type='type-name')) to categorize your nodes based on their function. For example, you might have a main coordinating node and several reporting worker nodes. You can set job_type to main for the main coordinating node and worker for the reporting worker nodes:

# Main coordinating node
with wandb.init(project="[PROJECT-NAME]", job_type="main", group="experiment_1") as run:
     # Training code

# Reporting worker nodes
with wandb.init(project="[PROJECT-NAME]", job_type="worker", group="experiment_1") as run:
     # Training code

Once you have set the job_type for your nodes, you can create saved views in your workspace to organize your runs. Click the action () menu and click Save as new view.

For example, you could create the following saved views:

• Default view: Filter out worker nodes to reduce noise.

  • Click Filter, then set Job Type to worker.
  • Shows only your reporting nodes.

• Debug view: Focus on worker nodes for troubleshooting.

  • Click Filter, then set Job Type == worker and set State to IN crashed.
  • Shows only worker nodes that have crashed or are in error states.

• All nodes view: See everything together.

  • No filter.
  • Useful for comprehensive monitoring.

To open a saved view, click Workspaces in the project sidebar, then click the menu. Workspaces appear at the top of the list and saved views appear at the bottom.

Track all processes to a single run

Parameters prefixed by x_ (such as x_label) are in public preview. Create a GitHub issue in the W&B repository to provide feedback.

Requirements

To track multiple processes to a single run, you must have:

• W&B Python SDK version v0.19.9 or newer.
• W&B Server v0.68 or newer.

In this approach you use a primary node and one or more worker nodes. Within the primary node you initialize a W&B run. For each worker node, initialize a run using the run ID used by the primary node. During training each worker node logs to the same run ID as the primary node. W&B aggregates metrics from all nodes and displays them in the W&B App UI.

Within the primary node, initialize a W&B run with wandb.init(). Pass in a wandb.Settings object to the settings parameter (wandb.init(settings=wandb.Settings()) with the following:

1. The mode parameter set to "shared" to enable shared mode.
2. A unique label for x_label. You use the value you specify for x_label to identify which node the data is coming from in logs and system metrics in the W&B App UI. If left unspecified, W&B creates a label for you using the hostname and a random hash.
3. Set the x_primary parameter to True to indicate that this is the primary node.
4. Optionally provide a list of GPU indexes ([0, 1, 2]) to x_stats_gpu_device_ids to specify which GPUs W&B tracks metrics for. If you don’t provide a list, W&B tracks metrics for all GPUs on the machine.

Make note of the run ID of the primary node. Each worker node requires the run ID of the primary node.

x_primary=True distinguishes a primary node from worker nodes. Primary nodes are the only nodes that upload files shared across nodes such as configuration files, telemetry, and more. Worker nodes don’t upload these files.

For each worker node, initialize a W&B run with wandb.init() and provide the following:

1. A wandb.Settings object to the settings parameter (wandb.init(settings=wandb.Settings()) with:
   • The mode parameter set to "shared" to enable shared mode.
   • A unique label for x_label. You use the value you specify for x_label to identify which node the data is coming from in logs and system metrics in the W&B App UI. If left unspecified, W&B creates a label for you using the hostname and a random hash.
   • Set the x_primary parameter to False to indicate that this is a worker node.
2. Pass the run ID used by the primary node to the id parameter.
3. Optionally set x_update_finish_state to False. This prevents non-primary nodes from updating the run’s state to finished prematurely, ensuring the run state remains consistent and managed by the primary node.

• Use the same entity and project for all nodes. This helps ensure the correct run ID is found.
• Consider defining an environment variable on each worker node to set the run ID of the primary node.

The following sample code demonstrates the high-level requirements for tracking multiple processes to a single run:

import wandb

entity = "[TEAM-ENTITY]"
project = "[PROJECT-NAME]"

# Initialize a run in the primary node
run = wandb.init(
    entity=entity,
    project=project,
	settings=wandb.Settings(
        x_label="rank_0", 
        mode="shared", 
        x_primary=True,
        x_stats_gpu_device_ids=[0, 1],  # (Optional) Only track metrics for GPU 0 and 1
        )
)

# Note the run ID of the primary node.
# Each worker node needs this run ID.
run_id = run.id

# Initialize a run in a worker node using the run ID of the primary node
run = wandb.init(
    entity=entity, # Use the same entity as the primary node
    project=project, # Use the same project as the primary node
	settings=wandb.Settings(x_label="rank_1", mode="shared", x_primary=False),
	id=run_id,
)

# Initialize a run in a worker node using the run ID of the primary node
run = wandb.init(
    entity=entity, # Use the same entity as the primary node
    project=project, # Use the same project as the primary node
	settings=wandb.Settings(x_label="rank_2", mode="shared", x_primary=False),
	id=run_id,
)

In a real-world example, each worker node might be on a separate machine.

See the Distributed Training with Shared Mode report for an end-to-end example on how to train a model on a multi-node and multi-GPU Kubernetes cluster in GKE.

To view console logs from multi-node processes in the project that the run logs to:

1. Navigate to the project that contains the run.
2. Click the Runs tab in the project sidebar.
3. Click the run you want to view.
4. Click the Logs tab in the project sidebar.

You can filter console logs based on the labels you provide for x_label in the UI search bar on the console log page. For example, the following image shows the available filter options when you provide values rank0, rank1, rank2, rank3, rank4, rank5, and rank6 to x_label.

See Console logs for more information.

W&B aggregates system metrics from all nodes and displays them in the W&B App UI. For example, the following image shows a sample dashboard with system metrics from multiple nodes. Each node possesses a unique label (rank_0, rank_1, rank_2) that you specify in the x_label parameter.

For more information about how to customize line plot panels, see Line plots.

Example use cases

The following code snippets demonstrate common scenarios for advanced distributed use cases. Use these patterns when your training setup involves spawned processes or when you need to share a single run across processes.

Spawn process

Use the wandb.setup() method in your main function if you initiate a run in a spawned process:

import multiprocessing as mp

def do_work(n):
    with wandb.init(config=dict(n=n)) as run:
        run.log(dict(this=n * n))

def main():
    wandb.setup()
    pool = mp.Pool(processes=4)
    pool.map(do_work, range(4))


if __name__ == "__main__":
    main()

Share a run

Pass a run object as an argument to share runs between processes:

def do_work(run):
    with wandb.init() as run:
        run.log(dict(this=1))

def main():
    run = wandb.init()
    p = mp.Process(target=do_work, kwargs=dict(run=run))
    p.start()
    p.join()
    run.finish()  # mark the run as finished


if __name__ == "__main__":
    main()

W&B can’t guarantee the logging order. The author of the script must handle synchronization.

Troubleshooting

The following sections describe two common issues you might encounter when using W&B with distributed training, along with guidance about enabling W&B Service to improve reliability:

• Becoming unresponsive at the beginning of training: A wandb process can stop responding if the wandb multiprocessing interferes with the multiprocessing from distributed training.
• Becoming unresponsive at the end of training: A training job might stop responding if the wandb process can’t determine when to exit. Call the wandb.Run.finish() API at the end of your Python script to signal W&B that the run finished. The wandb.Run.finish() API finishes uploading data and causes W&B to exit.

W&B recommends using wandb service command to improve the reliability of your distributed jobs. Both of the previous training issues commonly occur in versions of the W&B SDK where wandb service is unavailable.

Enable W&B Service

Depending on your version of the W&B SDK, you might already have W&B Service enabled by default.

W&B SDK 0.13.0 and later

W&B Service is enabled by default for versions of the W&B SDK 0.13.0 and later.

W&B SDK 0.12.5 and later

Modify your Python script to enable W&B Service for W&B SDK version 0.12.5 and later. Use the wandb.require() method and pass the string "service" within your main function:

if __name__ == "__main__":
    main()


def main():
    wandb.require("service")
    # rest-of-your-script-goes-here

W&B recommends that you upgrade to the latest version.

W&B SDK 0.12.4 and earlier

Set the WANDB_START_METHOD environment variable to "thread" to use multithreading instead if you use a W&B SDK version 0.12.4 or earlier.