> ## 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.

# Writing Custom Optimizers 

> Learn how to write a Cerebras-compliant custom optimizer.

Start by creating a subclass of [`cerebras.pytorch.optim.Optimizer`](https://training-api.cerebras.ai/en/latest/wsc/api/cerebras_pytorch/optim.html#cerebras.pytorch.optim.Optimizer "cerebras.pytorch.optim.Optimizer").

For example:

```python theme={null}
class CustomOptimizer(cstorch.optim.Optimizer):

    def __init__(self, params, ...):
        ...
        defaults = ...
        super().__init__(params, defaults, enable_global_step=...)

    ...

    def preinitialize(self):
        ...

    def step(self, closure=None):
        ...

    def state_names_to_sparsify(self):
        ...
```

As seen in the above example, similar to `torch.optim.Optimizer`, the base [`Optimizer`](https://training-api.cerebras.ai/en/latest/wsc/api/cerebras_pytorch/optim.html#cerebras.pytorch.optim.Optimizer "cerebras.pytorch.optim.Optimizer") class expects three arguments: the model parameters, the param group defaults, and the optional `enable_global_step` that will define a global step state variable for each parameter.

In addition, there are two abstract methods that must be overridden:

* [`preinitialize`](https://training-api.cerebras.ai/en/latest/wsc/api/cerebras_pytorch/optim.html#cerebras.pytorch.optim.Optimizer.preinitialize "cerebras.pytorch.optim.Optimizer.preinitialize")

This method is used to initialize any state variables that will be used by the optimizer. For example, [`SGD`](../../cs-torch/cerebras-pytorch-api/cerebras-pytorch-optim#class-cerebras-pytorch-optim-sgd-params-lr-momentum-0-dampening-0-weight-decay-0-nesterov-false-maximize-false "cerebras.pytorch.optim.SGD") defines its momentum buffers in its `preinitialize` method.

<Note>
  To remain Cerebras-compliant, no optimizer state variables may be initialized outside of the `preinitialize` method.
</Note>

For optimal performance, when initializing the state tensors that are filled with some constant value, use the [creation ops](https://training-api.cerebras.ai/en/latest/wsc/api/cerebras_pytorch/cstorch.html#creation-ops) that are available in the `cstorch` package to lazily initialize them.&#x20;

These ops will lazily initialize and fill the tensor, meaning that they take up very little memory and can be initialized much quicker than their `torch` counterparts when running on the cluster. Please see the source code for the optimizers in [cerebras.pytorch](https://training-api.cerebras.ai/en/latest/wsc/api/cerebras_pytorch/optim.html) for examples.

* [`step`](https://training-api.cerebras.ai/en/latest/wsc/api/cerebras_pytorch/optim.html#cerebras.pytorch.optim.Optimizer.step "cerebras.pytorch.optim.Optimizer.step")

This method is where the optimizer step is implemented. Note that due to the nature of lazy tensor tracing and execution, there may not be any Python level conditions or loops used to dynamically define the control flow. This means that only torch ops (such as `torch.where`) may be used.However, static structures are allowed. For example, a loop with a fixed number of iterations, or a Python conditional that doesn’t involve any torch tensors whose conditional involves only constant variables.

Once you’ve written your custom optimizer, as long as its available in the global scope, you can use it in ModelZoo in a similar way by setting the `optimizer_type` to be the name of your custom optimizer class in your params YAML file.