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

# Model

This page will cover how to pass a model into the [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html). The `model` is the main [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.4)") that all training and validation is run on. It is required by all [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html) instances.

## Prerequisites

Make sure to have read through [Trainer Overview](../../../model-zoo/trainer-overview) and [Trainer Configuration Overview](../../../model-zoo/trainer-configuration-overview) which provide the basic overview of how to run Model Zoo models. In this document, you will be using the tools and configurations outlined in those pages.

## Configure the Model

To set the model to train/validate using the [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html) use the `model` argument.

<CodeGroup>
  ```yaml YAML theme={null}

  All `model` subkeys are passed as arguments to the model class. The model class is decided by the `model_fn` in your run script.

  trainer:
    init:
      ...
      model:
        vocab_size: 1024
        max_position_embeddings: 1024
        ...
      ...
    ...

  ```

  ```python Python theme={null}
  The model can be passed as either:

  * a callable assumed to be a function that takes in no arguments returns a [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.4)")

  * a [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.4)") is passed and used as is


  from cerebras.modelzoo import Trainer
  from cerebras.modelzoo.models.nlp.gpt2.model import Gpt2Model

  trainer = Trainer(
      ...,
      model=lambda: Gpt2Model(
          vocab_size=1024,
          max_position_embeddings=1024,
          ...,
      ),
      ...,
  )
  ...
  ```
</CodeGroup>

<Note>
  If passing the `model` as a [`Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.4)") directly, it is optimal to first initalize the model inside of the Cerebras device context.

  For example:

  ```python theme={null}
  import cerebras.pytorch as cstorch
  from cerebras.modelzoo import Trainer
  from cerebras.modelzoo.models.nlp.gpt2.model import Gpt2Model

  # Initialize the Cerebras backend for efficient processing.
  backend = cstorch.backend("CSX")

  # Use the backend's device context manager for initializing the model.
  with backend.device:
      model = Gpt2Model(
          vocab_size=1024,
          max_position_embeddings=1024,
          ...,
      )

  # Compile the model using the Cerebras backend for optimized execution.
  trainer = Trainer(
      ...,
      backend=backend,
      model=model,
      ...,
  )
  ...
  ```

  This ensures that model parameters are automatically moved to the Cerebras device, optimizing memory usage and enhancing initialization speed. For more information, see [Efficient weight initialization](../../../cs-torch/writing-a-custom-training-loop/efficient-weight-initialization).
</Note>

## Conclusion

That covers specifying the model to train/validate with the [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html). You should now understand the various ways to configure the model and how the [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html) accepts a model.

## Further Reading

To learn more about how you can use the [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html) in some core workflows, you can check out:

* [Pretraining with Upstream Validation](../../../model-zoo/core-workflows/pretraining-with-upstream-validation)

To learn more about how you can extend the capabilities of the [`Trainer`](https://training-api.cerebras.ai/en/latest/wsc/Model-zoo/api/index.html) class, you can check out:

* [Defer Weight Initialization](../../../model-zoo/components/trainer-components/defer-weight-initialization)

* [Numeric Precision](../../../model-zoo/components/trainer-components/numeric-precision)

* [Train a model with weight sparsity](../../../model-zoo/tutorials/train-a-model-with-weight-sparsity)

* [Checkpointing](../../../model-zoo/components/trainer-components/checkpointing)

* [Customizing the Trainer with Callbacks](../../../model-zoo/components/trainer-components/customizing-the-trainer-with-callbacks)

* [Logging](../../../model-zoo/components/trainer-components/logging)