MosaicML Composer

State-of-the-art algorithms to train your neural networks

Composer is a library for training neural networks better, faster, and cheaper. It contains many state-of-the-art methods for accelerating neural network training and improving generalization, along with an optional Trainer API for composing many different enhancements.

This page shows you how to use W&B with Composer so you can track, visualize, and compare your training runs. W&B provides a wrapper for logging your ML experiments, but you don’t need to combine the two yourself: the Composer library incorporates W&B directly through the WandBLogger. This guide is for Composer users who want to log metrics, artifacts, and prediction samples to W&B.

Start logging to W&B

To start logging your Composer training runs to W&B, pass a WandBLogger instance to the Trainer:

from composer import Trainer
from composer.loggers import WandBLogger

trainer = Trainer(..., logger=WandBLogger())

Use Composer’s WandBLogger

The following sections describe how the WandBLogger integrates with Composer’s Trainer.

The Composer library uses the WandBLogger class in the Trainer to log metrics to W&B. Instantiate the logger and pass it to the Trainer:

wandb_logger = WandBLogger(project="gpt-5", log_artifacts=True)
trainer = Trainer(logger=wandb_logger)

Logger arguments

The following table describes the most common parameters you can use to customize how WandBLogger records your runs. For a full list and description, see the Composer documentation.

Parameter	Description
project	W&B project name (str, optional)
group	W&B group name (str, optional)
name	W&B run name. If not specified, uses the State.run_name (str, optional)
entity	W&B entity name, such as your username or W&B Team name (str, optional)
tags	W&B tags (List[str], optional)
log_artifacts	Whether to log checkpoints to W&B. Default: False (bool, optional)
rank_zero_only	Whether to log only on the rank-zero process. When logging artifacts, log on all ranks. W&B doesn’t store artifacts from ranks 1 and higher, which can discard relevant information. For example, when using Deepspeed ZeRO, you can’t restore from checkpoints without artifacts from all ranks. Default: True (bool, optional)
init_kwargs	Parameters to pass to wandb.init(), such as your W&B config. For the parameters that wandb.init() accepts, see wandb.init() parameters.

The following example shows a typical usage that passes run notes and a config dictionary through init_kwargs:

init_kwargs = {"notes":"Testing higher learning rate in this experiment", 
               "config":{"arch":"Llama",
                         "use_mixed_precision":True
                         }
               }

wandb_logger = WandBLogger(log_artifacts=True, init_kwargs=init_kwargs)

Log prediction samples

In addition to scalar metrics, you can log rich media such as model predictions to W&B for qualitative review. You can use Composer’s Callbacks system to control when you log to W&B through the WandBLogger. The following example logs a sample of the validation images and predictions:

import wandb
from composer import Callback, State, Logger

class LogPredictions(Callback):
    def __init__(self, num_samples=100, seed=1234):
        super().__init__()
        self.num_samples = num_samples
        self.data = []
        
    def eval_batch_end(self, state: State, logger: Logger):
        """Compute predictions per batch and stores them on self.data"""
        
        if state.timer.epoch == state.max_duration: #on last val epoch
            if len(self.data) < self.num_samples:
                n = self.num_samples
                x, y = state.batch_pair
                outputs = state.outputs.argmax(-1)
                data = [[wandb.Image(x_i), y_i, y_pred] for x_i, y_i, y_pred in list(zip(x[:n], y[:n], outputs[:n]))]
                self.data += data
            
    def eval_end(self, state: State, logger: Logger):
        with wandb.init() as run:
            "Create a wandb.Table and logs it"
            columns = ['image', 'ground truth', 'prediction']
            table = wandb.Table(columns=columns, data=self.data[:self.num_samples])
            run.log({'sample_table':table}, step=int(state.timer.batch))         
...

trainer = Trainer(
    ...
    loggers=[WandBLogger()],
    callbacks=[LogPredictions()]
)