> ## Documentation Index
> Fetch the complete documentation index at: https://training-docs.cerebras.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Customizing the Trainer with Callbacks

The [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html) class was designed to be easily extendable using [`Callback`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/generated/cerebras.modelzoo.trainer.callbacks.html#callback) classes. The `Trainer` exposes a number of hooks which can be overridden using a `Callback`.

On this page, you will learn about the basic `Callback` mechanism. By the end you should be able to write and use your own custom `Callback`.

## Prerequisites

Please ensure that you have read through the Cerebras Model Zoo Trainer Overview beforehand. The rest of this page assumes that you already have at least a cursory understanding of what the Cerebras Model Zoo Trainer is and how to use the python API.

## Callbacks

The callback mechanism is the backbone of the `Trainer`’s implementation. A lot of the heavy lifting in the Trainer is actually done by various [Core Callbacks](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/generated/cerebras.modelzoo.trainer.callbacks.html#core-callbacks).

In general, the `Callback` mechanism exposes a number of useful hooks that allow you to inject certain behaviour into the `Trainer`. These hooks include (but are not limited to)

* `setup`

* `on_{fit,train,validate}_{start,end}`

* `on_{train,validate}_batch_{start,end}`

* `on_{after,before}_{forward,backward}`

* `on_{after,before}_optimizer_{step,zero_grad}`

* `on_{after,before}_scheduler_step`

* `on_{save,load}_checkpoint`

* `on_after_save_checkpoint`

* `on_before_load_checkpoint`

The following pseudocode describes the structure of the [`fit`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html#cerebras.modelzoo.Trainer.fit) call and where the various hooks get called.

```python theme={null}
class Trainer:
    ...

    def fit(self, train_dataloader, val_dataloader, ckpt_path):

        on_before_load_checkpoint()
        load_checkpoint(ckpt_path)

        on_fit_start(...)

        for loop_idx in range(num_trains):
            on_train_start(...)

            for batch in train_dataloader:
                on_train_batch_start(...)

                on_before_forward(...)
                outputs = model(batch)
                on_after_forward(...)

                on_before_backward(...)
                outputs["loss"].backward()
                on_after_backward(...)

                on_before_optimizer_step(...)
                optimizer.step()
                on_after_optimizer_step(...)

                on_before_optimizer_zero_grad(...)
                optimizer.zero_grad()
                on_after_optimizer_zero_grad(...)

                for scheduler in schedulers:
                    on_before_scheduler_step(...)
                    scheduler.step()
                    on_after_scheduler_step(...)

                on_train_batch_end(...)

            on_train_end(...)

            for batch in val_dataloader:
                on_validate_batch_start(...)

                on_before_forward(...)
                outputs = model(batch)
                on_after_forward(...)

                on_validate_batch_end(...)


        on_fit_end(...)

```

For a comprehensive list of all supported hooks (as well as the arguments they accept), see the API docs for the `Callback` class.

## Pre-packaged Callbacks

There are many callbacks that come pre-packaged inside of the Model Zoo. See [Add-on Callbacks](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/generated/cerebras.modelzoo.trainer.callbacks.html#add-on-callbacks) for a complete list of all the callbacks available out-of-the-box in the Model Zoo

You can use any number of them to enhance the Trainer for your run.

For example,

<CodeGroup>
  ```python Python theme={null}
  from cerebras.modelzoo import Trainer
  from cerebras.modelzoo.trainer.callbacks import (
       ComputeNorm,
       CheckLoss,
       ...
  )

  trainer = Trainer(
      ...,
      callbacks=[
          ComputeNorm(),
          CheckLoss(),
          ...
      ],
      ...
  )
  ...

  ```

  ```yaml YAML theme={null}

  You can configure any pre-packaged callback inside a YAML configuration file as follows:

  trainer:
    init:
    ...
    callbacks:
    -  ComputeNorm:  {}
    -  CheckLoss:  {}
    ...

  ```
</CodeGroup>

## Global Callbacks

Any callback can be registered globally so that all `Trainer` instances know about it and will invoke that callback’s hooks.

There are two ways to globally register a callback. The first way is to treat the callback as a context manager. For example,

```python theme={null}
from cerebras.modelzoo import Trainer
from cerebras.modelzoo.trainer.callbacks import CheckLoss

with CheckLoss():
    trainer = Trainer(...)
    trainer.fit(...)

```

In the above example, while within the `CheckLoss`’s context, all trainer `fit` calls inside the context will check the loss values that come out of the model.

The other way to register a callback is to call :py:function:`~cerebras.modelzoo.trainer.callbacks.register\_global\_callback`.

For example,

```python theme={null}
from cerebras.modelzoo import Trainer
from cerebras.modelzoo.trainer.callbacks import CheckLoss, register_global_callback

callback = CheckLoss()
handle = register_global_callback(callback)

trainer = Trainer(...)
trainer.fit(...)

handle.remove()

```

In the above example, all trainer `fit` calls inside the context will check the loss values that come out of the model.

:py:function:`~cerebras.modelzoo.trainer.callbacks.register\_global\_callback` returns a removeable handle object that can be used to remove the added callback by calling `handle.remove()`

## Callback Ordering

The `Trainer` is comprised of many different callbacks that all serve to enhance its functionality\`.

All of these callbacks share common hooks. These hooks must be called in a specific order. The order in which callbacks are invoked is as follows:

1. Core Callbacks: The callbacks that implement the most fundamental behaviour of the `Trainer` get called first.

2. User-defined callbacks: The callbacks that are passed into the `callbacks` argument of the `Trainer`’s constructor are called next.

3. Global callbacks: Finally, the callbacks that are registered globally are called.

For example,

```python theme={null}
from cerebras.modelzoo import Trainer
from cerebras.modelzoo.trainer.callbacks import (
    CheckLoss,
    ComputeNorm,
    TrainingLoop,
)

trainer = Trainer(
    ...,
    loop=TrainingLoop(num_steps=1),
    callbacks=[ComputeNorm()],
    ...
)

with CheckLoss():
    trainer.fit(...)

```

Let’s consider the `on_fit_start` hook. Between the three callbacks that are highlighed in the above example, the order that the callbacks’s `on_fit_start` hook is invoked is as follows:

1. `TrainingLoop.on_fit_start`: As `TrainingLoop` is a core callback.

2. `ComputeNorm.on_fit_start`: As `ComputeNorm` was passed into the Trainer’s constructor.

3. `CheckLoss.on_fit_start`: As it is a globally registered callback.

## Writing a Custom Callback

To write your own custom callback class, all you need to do is inherit from the base `Callback` class and override the hooks that you need.

For example, let’s implement a simple callback that scales the loss value by some constant value before we call `loss.backward()`

```python theme={null}
from cerebras.modelzoo.trainer.callbacks import Callback

class ScaleLoss(Callback):
    def __init__(self, value):
        self.value = value

    def on_before_backward(self, trainer, model, outputs):
        outputs["loss"] *= value

```

That is all there is to it. This callback can now be used inside the `Trainer` as follows:

<CodeGroup>
  ```python Python theme={null}

  from cerebras.modelzoo import Trainer

  trainer = Trainer(
      ...,
      callbacks=[ScaleLoss(value=0.95)],
      ...
  )
  ...

  ```

  ```yaml YAML theme={null}
  As long as the callback class exists in the Python global namespace, you can add any custom callback to a YAML configuration file in exactly the same way you would any other pre-packaged callback.

  trainer:
    init:
    ...
    callbacks:
    -  ScaleLoss:
    value:  0.95
    ...

  ```
</CodeGroup>

In order for the callback class to exist in the Python global namespace, the Python interpreter must have seen it at some point. Implementing your custom callback in the same file as the model class ensures that the callback is seen by the Python interpreter and loaded into the Python global namespace.