Model Zoo Config Classes

When using the Cerebras Model Zoo, working with model parameters and configurations is a common and important task. To make it easy for you to define, manage, and validate model configurations, Model Zoo uses Config classes built on Pydantic’s BaseModel framework.

Config classes offer a robust and structured approach to managing model configurations. These classes act as detailed “schemas” that store essential model settings and enforce data type and value constraints, providing a much-needed layer of validation.

Furthermore, Config classes allow developers to manage and tweak settings for each model without the need for complex code alterations. This approach fosters a more structured and adaptable environment for experimentation and fine-tuning.

​

Key Advantages of Config Classes

• Validation: By defining expected data types and constraints using Pydantic validators, Config classes significantly reduce the risk of common errors such as typos, incorrect data types, or missing fields.

• Improved Documentation: With the structured nature of Config classes, documentation will be automatically generated from docstrings present in the class.

• Enhanced Manageability: As models evolve to become more sophisticated, Config classes provide a clear and organized way to manage configuration parameters within the Python code.

​

Config Class Hierarchy

The Config classes follow a clear inheritance hierarchy:

1. BaseConfig: The root class that inherits from Pydantic’s BaseModel, providing core validation functionality

2. ModelConfig: Inherits from BaseConfig, specifically for model configurations

3. DataConfig: Inherits from BaseConfig, specifically for data processor configurations

These base classes can be imported from the cerebras.modelzoo.config module:

from cerebras.modelzoo.config import BaseConfig, ModelConfig, DataConfig

​

BaseConfig

BaseConfig is a subclass of Pydantic’s BaseModel class with specific constraints. Config classes in Model Zoo are immutable by default, which means that once a config object is created, its attributes cannot be modified. This immutability helps ensure configuration stability during model execution.

However, there are cases where you need to set some attributes based on other values during initialization. For these situations, the only place where attribute assignments are allowed is within the post_init method. For example:

from cerebras.modelzoo.config import BaseConfig

class SomeConfig(BaseConfig):
    a: int = ...  # ... in pydantic means field is required
    b: str = "undefined"
    
    def post_init(self, context):
        # Only place where attribute assignments are allowed
        if self.a > 10:
            self.b = "greater than 10"
        else:
            self.b = "less than or equal to 10"

Alternatively, you can use Pydantic’s field_validators to change the value of a field:

import os
from pydantic import field_validator
from cerebras.modelzoo.config import BaseConfig

class SomeConfig(BaseConfig):
    some_path: str
    
    @field_validator("some_path")
    @classmethod
    def change_some_path(cls, some_path):
         return os.path.realpath(some_path)

​

ModelConfig

ModelConfig is a subclass of BaseConfig with the additional requirement that any config classes inheriting from it must specify a name field. This field’s type must be defined as Literal["name_of_the_model"]. The value of this name field is what the system uses to determine which model to initialize.

from cerebras.modelzoo.config import ModelConfig
from typing import Literal

class MyModelConfig(ModelConfig):
    name: Literal["my_model_name"]  # Required field
    # Other model-specific fields

​

DataConfig

DataConfig is a subclass of BaseConfig with the additional requirement that any config classes inheriting from it must specify a data_processor field. This field’s type must be defined as Literal["name_of_the_data_processor"], which ensures the field exactly matches the specific data processor identifier the system should initialize.

from cerebras.modelzoo.config import DataConfig
from typing import Literal

class MyDataProcessorConfig(DataConfig):
    data_processor: Literal["name_of_the_data_processor"]  # Required field
    # Other data processor-specific fields

​

Validation Mechanisms

Config classes in Model Zoo provide robust validation capabilities through Pydantic’s validation system. There are two main approaches to validation:

• Field-level validation

• Model-level validation

​

Field-Level Validation

Field validators check individual parameters without considering other fields. Use field-level validation when a parameter’s constraints are independent of other configuration values.

from pydantic import field_validator
from cerebras.modelzoo.config import BaseConfig

class SomeConfig(BaseConfig):
    some_path: str
  
    @field_validator("some_path")
    @classmethod
    def check_some_path(cls, some_path):
         if os.path.isabs(some_path):
             raise ValueError("Only relative paths allowed")
         return some_path

​

Model-Level Validation

Use model validators when parameters have interdependencies or when validation requires checking multiple fields together.

from pydantic import model_validator
from cerebras.modelzoo.config import BaseConfig

class SomeConfig(BaseConfig):
    should_be_relative: bool
    some_path: str
    @model_validator(mode="after")
    def check_some_path(self):
         if self.should_be_relative and os.path.isabs(self.some_path):
             raise ValueError("Only relative paths allowed")
         return self

​

Config Validation Tool

Model Zoo provides a CLI tool to validate your configuration files before running your model. This helps catch configuration errors early in development.

To validate a configuration file using the CLI:

cszoo config validate /path/to/params.yaml

​

Further Reading

• For detailed information about validation mechanisms, refer to the Pydantic validation documentation.

• For Model Zoo-specific implementations, review the examples in the Model Zoo repository.