composer.Algorithm

Algorithms are implemented in both a standalone functional form (see composer.functional) and as subclasses of Algorithm for integration in the MosaicML Trainer. This section describes the latter form.

For ease of composability, algorithms in our Trainer are based on the two-way callbacks concept from Howard et al., 2020. Each algorithm implements two methods:

  • Algorithm.match(): returns True if the algorithm should be run given the current State and Event.

  • Algorithm.apply(): performs an in-place modification of the given State

For example, a simple algorithm that shortens training:

from composer import Algorithm, State, Event, Logger

class ShortenTraining(Algorithm):

    def match(self, state: State, event: Event, logger: Logger) -> bool:
        return event == Event.TRAINING_START

    def apply(self, state: State, event: Event, logger: Logger):
        state.max_epochs /= 2  # cut training time in half

For a complete list of algorithms, see composer.algorithms.

For reference, available events include:

Name

Description

INIT

Immediately after model initialization, and before creation of optimizers and schedulers. Model surgery typically occurs here.

TRAINING_START

Start of training. For multi-GPU training, runs after the DDP process fork.

EPOCH_START, EPOCH_END

Start and end of an Epoch.

BATCH_START, BATCH_END

Start and end of a batch, inclusive of the optimizer step and any gradient scaling.

AFTER_DATALOADER

Immediately after the dataloader is called. Typically used for on-GPU dataloader transforms.

BEFORE_TRAIN_BATCH, AFTER_TRAIN_BATCH

Before and after the forward-loss-backward computation for a training batch. When using gradient_accumulation, these are still called only once.

BEFORE_FORWARD, AFTER_FORWARD

Before and after the call to model.forward()

BEFORE_LOSS, AFTER_LOSS

Before and after the loss computation.

BEFORE_BACKWARD, AFTER_BACKWARD

Before and after the backward pass.

TRAINING_END

End of training.

EVAL_START, EVAL_END

Start and end of evaluation through the validation dataset.

EVAL_BATCH_START, EVAL_BATCH_END

Before and after the call to model.validate(batch)

EVAL_BEFORE_FORWARD, EVAL_AFTER_FORWARD

Before and after the call to model.validate(batch)

For more information about events, see composer.Event.