Writing a custom training loop#

Overview#

Our typical workflow involves using a training script provided in the Cerebras Model Zoo. However, if that training loop is insufficient for your model needs, you may write your own training loop using the Cerebras PyTorch API.

Proceed with the following steps to learn how to write a custom training loop for a simple, fully connected model for training on the MNIST dataset.

Note, the following steps will only take you through the absolute minimum code required to run a simple, small model on the Cerebras Wafer Scale Cluster. To extend the script to feature things like learning rate scheduling, gradient scaling, etc. please continue to the further reading section to learn more about these topics.

Prerequisites#

You have installed the cerebras.pytorch package in your environment.

Validate the package installation#

To check whether the cerebras.pytorch package is installed correctly, issue the following command:

import cerebras.pytorch as cstorch

Note

From here on, we will be using cstorch as the alias for cerebras.pytorch

Define your model#

When using the Cerebras PyTorch API, you can define your model in the same way you would in a Vanilla PyTorch workflow:

import torch
import torch.nn.functional as F


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(784, 256)
        self.fc2 = torch.nn.Linear(256, 10)

    def forward(self, inputs):
        inputs = torch.flatten(inputs, 1)
        outputs = F.relu(self.fc1(inputs))
        return F.relu(self.fc2(outputs))


model = Model()

Note

Weight initialization for large models can cause out-of-memory errors. Not only that, but initializing extremely large models eagerly can be very slow. See the page on Efficient weight initialization to see how to work around this issue.

Compile your model#

Once the model has been instantiated, compile the model by calling the cerebras.pytorch.compile, e.g.

compiled_model = cstorch.compile(model, backend="CSX")

You must pass in the backend you wish to compile the model with. You can simply pass in the type of backend if you wish to use all default arguments, or you can instantiate a backend object using cerebras.pytorch.backend to customize it, e.g.

backend = cstorch.backend("CSX", artifact_dir="./artifact_dir")
compiled_model = cstorch.compile(model, backend)

Note

The call to cstorch.compile doesn’t actually compile the model. Similar to torch.compile it only prepares the model for compilation. Compilation only happens after the first iteration once the input shapes are known.

Optimize Model Parameters#

To optimize model parameters using the Cerebras Wafer-Scale cluster, you must use a Cerebras-compliant optimizer. There are exact drop-in replacements for all commonly used optimizers available in cerebras.pytorch.optim, e.g.

optimizer = cstorch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

Note

For convenience, we also include a configuration helper function configure_optimizer.

If you are interested in writing your own Cerebras custom-compliant optimizer, see the page on Writing custom optimizers

DataLoaders#

To send data to the Wafer-Scale cluster, you must wrap your PyTorch dataloader with cerebras.pytorch.utils.data.DataLoader, e.g.

def get_torch_dataloader(batch_size, train):
    from torchvision import datasets, transforms

    train_dataset = datasets.MNIST(
        "./data",
        train=train,
        download=True,
        transform=transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,)),
            ]
        ),
        target_transform=transforms.Lambda(
            lambda x: torch.as_tensor(x, dtype=torch.int32)
        ),
    )

    return torch.utils.data.DataLoader(
        train_dataset, batch_size=batch_size, shuffle=True
    )


training_dataloader = cstorch.utils.data.DataLoader(
    get_torch_dataloader, batch_size=64, train=True
)

The Cerebras PyTorch dataloader takes in some callable that returns a PyTorch dataloader. It must be done this way so that every single worker can create their own PyTorch dataloader instance to maximize distributed parallelism.

Define the Training Step#

To run a single training iteration on the Cerebras Wafer-Scale cluster, we must first, capture everything that is intended to run on the cluster. To do this, define a function which contains everything that happens in a single training iteration, and decorate it using cerebras.pytorch.trace.

For example:

loss_fn = torch.nn.CrossEntropyLoss()

@cstorch.trace
def training_step(inputs, targets):
    outputs = compiled_model(inputs)
    loss = loss_fn(outputs, targets)

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

    return loss

This function gets traced and sent to the cluster for compilation and execution.

Define an Execution#

To program an execution run on the Cerebras Wafer-Scale cluster, you must define an instance of the cerebras.pytorch.utils.data.DataExecutor, e.g.

train_executor = cstorch.utils.data.DataExecutor(
    training_dataloader,
    num_steps=100,
    checkpoint_steps=50,
)

It takes in the Cerebras PyTorch dataloader that will be used during the run, the total number of steps to run for, as well as the interval at which checkpoints will be taken.

Configuring the Cerebras Wafer Scale Cluster#

To configure the Cerebras Wafer-Scale cluster, construct a CSConfig object:

cs_config = cstorch.utils.CSConfig(
    mgmt_address=mgmt_address,
    transfer_processes=1,
    max_wgt_servers=1,
    max_act_per_csx=1,
    num_workers_per_csx=1,
)

which can be passed to the cerebras.pytorch.utils.data.DataExecutor object

train_executor = cstorch.utils.data.DataExecutor(
    training_dataloader,
    num_steps=100,
    checkpoint_steps=50,
    cs_config=cs_config,
)

See the class documentation for CSConfig for all the options configurable.

Note

Most options have reasonable defaults and do not need to be changed.

Train your model#

Once the above is defined, you can iterate through the executor to train your model.

@cstorch.step_closure
def print_loss(mode, loss: torch.Tensor, step: int):
    print(f"{mode} Loss {step}: {loss.item()}")


@cstorch.checkpoint_closure
def save_checkpoint(step):
    cstorch.save(
        {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
        },
        f"checkpoint_{step}.mdl",
    )

global_step = 0

for inputs, targets in train_executor:
    loss = training_step(inputs, targets)
    print_loss("Training", loss, global_step)
    global_step += 1
    save_checkpoint(global_step)

Note

Notice how the loss was passed into a function decorated by step_closure. This is required to retrieve the loss value from the Cerebras Wafer Scale Cluster before it can be used. Please see the page on step closures for more details.

Note

Also, notice how checkpoints are saved inside a function decorated by checkpoint_closure. This is required to retrieve the model weights and optimizer state back from the Cerebras Wafer Scale Cluster before it can be saved. Please see the page on saving checkpoints

Putting it all together#

Combining all of the above steps, we can create a super minimal training script for a simple, fully connected model training on the MNIST dataset:

# Import the Cerebras PyTorch module
import cerebras.pytorch as cstorch

# Define a model
import torch
import torch.nn.functional as F


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(784, 256)
        self.fc2 = torch.nn.Linear(256, 10)

    def forward(self, inputs):
        inputs = torch.flatten(inputs, 1)
        outputs = F.relu(self.fc1(inputs))
        return F.relu(self.fc2(outputs))


model = Model()

# Compile the model
compiled_model = cstorch.compile(model, backend="CSX")

# Define an optimizer
optimizer = cstorch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)


# Define a data loader
def get_torch_dataloader(batch_size, train):
    from torchvision import datasets, transforms

    train_dataset = datasets.MNIST(
        "./data",
        train=train,
        download=True,
        transform=transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,)),
            ]
        ),
        target_transform=transforms.Lambda(
            lambda x: torch.as_tensor(x, dtype=torch.int32)
        ),
    )

    return torch.utils.data.DataLoader(
        train_dataset, batch_size=batch_size, shuffle=True
    )


training_dataloader = cstorch.utils.data.DataLoader(
    get_torch_dataloader, batch_size=64, train=True
)

# Define the training step
loss_fn = torch.nn.CrossEntropyLoss()

@cstorch.trace
def training_step(inputs, targets):
    outputs = compiled_model(inputs)
    loss = loss_fn(outputs, targets)

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

    return loss


@cstorch.step_closure
def print_loss(loss: torch.Tensor, step: int):
    print(f"Train Loss {step}: {loss.item()}")


@cstorch.checkpoint_closure
def save_checkpoint(step):
    cstorch.save(
        {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
        },
        f"checkpoint_{step}.mdl",
    )


cs_config = cstorch.utils.CSConfig(
    mgmt_address=mgmt_address,
    transfer_processes=1,
    max_wgt_servers=1,
    max_act_per_csx=1,
    num_workers_per_csx=1,
)

global_step = 0

train_executor = cstorch.utils.data.DataExecutor(
    training_dataloader,
    num_steps=100,
    checkpoint_steps=50,
    cs_config=cs_config,
)
model.train()
for inputs, targets in train_executor:
    loss = training_step(inputs, targets)
    print_loss(loss, global_step)
    global_step += 1
    save_checkpoint(global_step)

Writing a custom training loop#

Overview#

Prerequisites#

Validate the package installation#

Define your model#

Compile your model#

Optimize Model Parameters#

DataLoaders#

Define the Training Step#

Define an Execution#

Configuring the Cerebras Wafer Scale Cluster#

Train your model#

Putting it all together#

Further reading#