The Model Zoo Trainer helps streamline the process of training large AI models by taking care of complex setup steps, like manual parallelism. It provides a fast, scalable way to train and validate models.

While you don’t need to use the Trainer to run the reference models included in the Model Zoo, it’s recommended—especially if you want to take advantage of Cerebras hardware performance.

This page introduces the Trainer class and gives you a basic understanding of how to use it.

​
Prerequisites

Please ensure that you have installed the Cerebras Model Zoo package by going through the installation guide.

Optionally, you can also read through the basic Cerebras PyTorch guide to first gain an understanding of the underlying API that underpins the Trainer class.

​
Basic Usage

The Trainer class can be imported and used as follows:


import cerebras.pytorch as cstorch
from cerebras.modelzoo import Trainer

# Any torch.nn.Module
model: torch.nn.Module = torch.nn.Linear(10, 10)
# Any Cerebras compliant optimizer
optimizer: cstorch.optim.Optimizer = cstorch.optim.SGD(
    model.parameters(), lr=0.01, momentum=0.9
)

trainer = Trainer(
    device="CSX",  # The device to run on
    model_dir="./model_dir",  # The directory at which to store artifacts
    model=model,
    optimizer=optimizer,
)
# Train the model over a single epoch of the train dataloader
# and then run validation over a single epoch of the val dataloader
trainer.fit(train_dataloader, val_dataloader)

As can be seen in the above example, at a minimum the Trainer class takes in the following:

Learn more about these parameters in our Trainer Configuration guide.

  • device: The device to run training/validation on.

  • model_dir: The directory at which to store model related artifacts (e.g. model checkpoints).

  • model: The torch.nn.Module instance that we are training/validating.

  • optimizer: Optionally, a cerebras.pytorch.optim.Optimizer instance can be passed in to optimize the model weights during the training phase.

At a minimum, the call to fit takes in the following:

The default behaviour of this minimally configured run is to train the model over a single epoch of the train dataloader and then run validation over a single epoch of the val dataloader.

There you have it! With this small sample of code, you can begin training your very first model using the Cerebras Model Zoo Trainer!

You can pause here to go try it out for yourself, or continue reading to learn more about how to more finely configure the Trainer to fit your needs.

​
Configuring the Training loop

As mentioned above, if both a train_dataloader and val_dataloader are provided to the fit call, the default behaviour is to run a single epoch of training followed by a single epoch of validation.

This behaviour can be configured by passing in a TrainingLoop instance to the Trainer as follows:

from cerebras.modelzoo.trainer.callbacks import TrainingLoop

trainer = Trainer(
      ...,
      loop=TrainingLoop(
          num_steps=1000,
          eval_steps=100,
          eval_frequency=100,
      ),
)
trainer.fit(train_dataloader, val_dataloader)

In this above example,

  • num_steps represents the total number of batches to train for. If num_steps exceeds the number of available batches in the train dataloader, the dataloader is automatically repeated to be able to run training for num_steps.

  • eval_steps represents the number of steps to run validation for every time we run validation. Similar to training, if eval_steps exceeds the number of available batches in the val dataloader, the dataloader is automatically repeated. Although, typically validation is never run for more than a single epoch. So, it is advised to set eval_steps to be less than the length of the validation dataloader. Otherwise, the validation metrics may be incorrect.

  • eval_frequency represents how often validation is run during training. In the above example, validation is run every 100 steps of training. That is to say, throughout the 1000 steps of training, validation is run 10 times. Regardless of the value of eval_frequency, if eval_frequency is greater than zero, we always run validation at the end of training.

​
Checkpointing

The Trainer can be further configured to save checkpoints at regular intervals by passing in a Checkpoint instance as follows:

from cerebras.modelzoo.trainer.callbacks import Checkpoint

trainer = Trainer(
      ...,
      model_dir="./model_dir",
      checkpoint=Checkpoint(steps=100),
)
trainer.fit(train_dataloader, val_dataloader)

In the above example, a checkpoint is saved every 100 steps of training. A checkpoint is also saved at the end of training regardless of whether or not num_steps is a multiple of the checkpoint steps.

The checkpoints are saved in the model_dir directory that was passed to the Trainer.

model_dir/
├── checkpoint_100.mdl
├── checkpoint_200.mdl
├── checkpoint_300.mdl
└── ...

This checkpoint is meant for resuming training from the same point in the future. As such, it will contain the model weights, optimizer state, and any other state that is necessary to resume training. Please see Selective Checkpoint State Saving for examples of how to configure what state is saved into the checkpoint.

A saved checkpoint can be loaded again in the future by specifying the ckpt_path argument to the call to fit. For example,

trainer = Trainer(...)
trainer.fit(
    train_dataloader,
    val_dataloader,
    ckpt_path="/path/to/checkpoint",
)

The above code will load the checkpoint at that path before starting training.

If a ckpt_path is not provided, but a checkpoint is found inside the model_dir, then Trainer “cerebras.modelzoo.Trainer”) will automatically load the latest checkpoint found in the model_dir.

To learn more about how to configure checkpointing behavior using the Trainer, see Checkpointing.

​
What’s next?

To learn about how to specify a schedule for learning rates, please see Optimizer and Scheduler.

To learn about how you can configure a Trainer instance using a YAML configuration file, you can check out:

  • Trainer YAML Overview

To learn more about how you can use the Trainer in some core workflows, you can check out:

To learn more about how you can extend the capabilities of the Trainer class, you can check out:

To learn more about what the Trainer class outputs during the run, you can check out: