Validated Dataclasses for Model and Experiment Configuration

Why configuration deserves a real model

A small script can get away with loose variables:

learning_rate = 3e-4
batch_size = 8
num_heads = 12
hidden_size = 768

As an experiment grows, this becomes fragile. Settings are scattered, names drift, invalid combinations are easy to create, and reproducing a run means hunting through source files, command-line history, and notebooks.

A configuration object establishes one named, inspectable contract:

ExperimentConfig(
    project_name="vlm-retrieval-baseline",
    model=ModelConfig(...),
    training=TrainingConfig(...),
)

For an applied AI system, configuration should answer three questions:

1. What settings define this run?
2. Are those settings internally valid?
3. Can the run be reproduced and compared later?

Dataclasses help with the first and third questions. Validation handles the second.

A standard dataclass generates useful methods such as an initializer, a readable representation, and equality comparison from annotated fields. But it is important to retain the distinction from the previous lesson:

Type hints describe what a value is expected to be; validation decides whether a supplied value is acceptable.

For instance, num_heads: int rules out a string under mypy, but it does not express that num_heads must be positive or that hidden_size % num_heads == 0.

Dataclass mechanics for configurations

Create src/vlm_lab/config.py. We will use three small immutable configuration classes:

• ModelConfig describes a decoder-style transformer’s structural choices.
• TrainingConfig describes optimization choices.
• ExperimentConfig groups them into a single run specification.

Before writing the implementation, study the relevant standard-library behavior.

Here is the implementation:

"""Typed, validated configuration for VLM experiments."""

from dataclasses import dataclass
from math import isfinite
from typing import Literal, TypeAlias


DType: TypeAlias = Literal["float32", "bfloat16"]


def _require_nonempty_string(name: str, value: object) -> None:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(
            f"{name} must be a non-empty string; got {value!r}"
        )


def _require_int_at_least(
    name: str,
    value: object,
    minimum: int,
) -> None:
    if (
        isinstance(value, bool)
        or not isinstance(value, int)
        or value < minimum
    ):
        raise ValueError(
            f"{name} must be an integer >= {minimum}; got {value!r}"
        )


def _require_finite_real_at_least(
    name: str,
    value: object,
    minimum: float,
) -> None:
    if (
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not isfinite(value)
        or value < minimum
    ):
        raise ValueError(
            f"{name} must be a finite number >= {minimum}; got {value!r}"
        )


@dataclass(frozen=True, slots=True, kw_only=True)
class ModelConfig:
    """Architecture and numeric-format settings for one model."""

    model_name: str
    vocab_size: int
    hidden_size: int = 768
    num_layers: int = 12
    num_heads: int = 12
    max_seq_len: int = 512
    dropout: float = 0.0
    dtype: DType = "bfloat16"

    def __post_init__(self) -> None:
        _require_nonempty_string("model_name", self.model_name)
        _require_int_at_least("vocab_size", self.vocab_size, 2)
        _require_int_at_least("hidden_size", self.hidden_size, 1)
        _require_int_at_least("num_layers", self.num_layers, 1)
        _require_int_at_least("num_heads", self.num_heads, 1)
        _require_int_at_least("max_seq_len", self.max_seq_len, 2)

        if self.hidden_size % self.num_heads != 0:
            raise ValueError(
                "hidden_size must be divisible by num_heads; "
                f"got hidden_size={self.hidden_size} and "
                f"num_heads={self.num_heads}"
            )

        if (
            not isinstance(self.dropout, (int, float))
            or isinstance(self.dropout, bool)
            or not isfinite(self.dropout)
            or not 0.0 <= self.dropout < 1.0
        ):
            raise ValueError(
                "dropout must be a finite number in [0.0, 1.0); "
                f"got {self.dropout!r}"
            )

        if self.dtype not in ("float32", "bfloat16"):
            raise ValueError(
                "dtype must be 'float32' or 'bfloat16'; "
                f"got {self.dtype!r}"
            )


@dataclass(frozen=True, slots=True, kw_only=True)
class TrainingConfig:
    """Optimization settings for one training run."""

    learning_rate: float = 3e-4
    batch_size: int = 8
    epochs: int = 3
    seed: int = 1337
    weight_decay: float = 0.01
    max_grad_norm: float | None = 1.0

    def __post_init__(self) -> None:
        _require_finite_real_at_least(
            "learning_rate",
            self.learning_rate,
            minimum=0.0,
        )
        if self.learning_rate == 0.0:
            raise ValueError("learning_rate must be greater than zero")

        _require_int_at_least("batch_size", self.batch_size, 1)
        _require_int_at_least("epochs", self.epochs, 1)
        _require_int_at_least("seed", self.seed, 0)
        _require_finite_real_at_least(
            "weight_decay",
            self.weight_decay,
            minimum=0.0,
        )

        if self.max_grad_norm is not None:
            _require_finite_real_at_least(
                "max_grad_norm",
                self.max_grad_norm,
                minimum=0.0,
            )
            if self.max_grad_norm == 0.0:
                raise ValueError(
                    "max_grad_norm must be positive or None"
                )


@dataclass(frozen=True, slots=True, kw_only=True)
class ExperimentConfig:
    """Complete configuration for a named experiment."""

    project_name: str
    model: ModelConfig
    training: TrainingConfig
    tags: tuple[str, ...] = ()

    def __post_init__(self) -> None:
        _require_nonempty_string("project_name", self.project_name)

        if not isinstance(self.model, ModelConfig):
            raise ValueError("model must be a ModelConfig instance")

        if not isinstance(self.training, TrainingConfig):
            raise ValueError(
                "training must be a TrainingConfig instance"
            )

        if not isinstance(self.tags, tuple):
            raise ValueError("tags must be a tuple of strings")

        if any(
            not isinstance(tag, str) or not tag.strip()
            for tag in self.tags
        ):
            raise ValueError(
                "every tag must be a non-empty string"
            )

        if len(set(self.tags)) != len(self.tags):
            raise ValueError("tags must not contain duplicates")

The code is longer than a collection of loose constants because it encodes decisions that otherwise remain implicit. A model with 768 hidden features and 12 heads has head dimension $768/12=64$. With 10 heads, there is no integral partition into equal attention-head subspaces; the configuration is structurally invalid before any tensor is allocated.

Why these decorator options?

@dataclass(frozen=True, slots=True, kw_only=True)

Each option has a distinct role:

Freezing is not deep immutability. If a frozen dataclass contains a mutable list, Python still allows mutation inside that list. That is why the main configuration uses tuple[str, ...] for tags: a tuple reflects the intention that experiment labels are fixed.

Constructing a valid experiment

Add the following exports to src/vlm_lab/__init__.py, preserving your existing build_prompt export:

from .config import ExperimentConfig, ModelConfig, TrainingConfig
from .prompts import build_prompt

__all__ = [
    "ExperimentConfig",
    "ModelConfig",
    "TrainingConfig",
    "build_prompt",
]

From the project root, open a Python REPL with your virtual environment activated and construct a configuration:

from vlm_lab import ExperimentConfig, ModelConfig, TrainingConfig

config = ExperimentConfig(
    project_name="vlm-retrieval-baseline",
    model=ModelConfig(
        model_name="toy-vlm-decoder",
        vocab_size=32_000,
        hidden_size=768,
        num_layers=12,
        num_heads=12,
        max_seq_len=512,
        dropout=0.1,
        dtype="bfloat16",
    ),
    training=TrainingConfig(
        learning_rate=3e-4,
        batch_size=8,
        epochs=3,
        seed=1337,
        weight_decay=0.01,
    ),
    tags=("baseline", "retrieval"),
)

print(config)

The generated representation is useful in logs and experiment records. Because dataclasses also generate value-based equality by default, two independently constructed configurations with identical fields compare equal:

same_config = ExperimentConfig(
    project_name="vlm-retrieval-baseline",
    model=ModelConfig(
        model_name="toy-vlm-decoder",
        vocab_size=32_000,
        dropout=0.1,
    ),
    training=TrainingConfig(),
    tags=("baseline", "retrieval"),
)

print(config == same_config)

This example prints True only if every field, including defaults, matches. That property becomes useful later when you want to check whether a checkpoint was resumed with the same configuration.

Now try an invalid architecture:

ModelConfig(
    model_name="invalid-transformer",
    vocab_size=32_000,
    hidden_size=768,
    num_heads=10,
)

Construction immediately raises a ValueError explaining the violated relationship. This is the desired failure mode: clear, local, and before GPU memory is allocated or a long run begins.

Finally, re-run static checking:

python -m mypy src

Mypy and runtime validation complement each other:

• mypy catches a source-code mismatch such as batch_size="eight" before execution.
• __post_init__ rejects invalid numeric values and incompatible combinations during construction.
• Neither one verifies whether a model name exists in a remote registry or whether the available GPU has enough memory. Those are environment-dependent checks, which belong in the component that loads or runs the model.

Defaults, factories, and configuration mutability

Direct immutable defaults are safe:

tags: tuple[str, ...] = ()

But this is wrong in ordinary Python classes and prohibited by modern dataclasses:

# Do not write this.
overrides: list[str] = []

A mutable list created once would be shared by every configuration instance. If a future sweep configuration genuinely needs mutable state, use a factory:

from dataclasses import dataclass, field


@dataclass
class SweepConfig:
    overrides: list[str] = field(default_factory=list)

default_factory=list calls list() separately for each new SweepConfig. No instance shares another’s default list.

For durable model and experiment configuration, favor immutable fields where practical:

• use tuples instead of lists for fixed sequences;
• use nested frozen dataclasses rather than mutable dictionaries;
• create a new configuration rather than modifying an existing one.

That approach makes a configuration closer to a versioned experiment specification than a bag of settings that can drift during a run.

When standard dataclasses are not enough

The standard-library approach is a strong default when configuration is constructed inside typed Python code. It is lightweight, transparent, and gives you complete control over validation messages and cross-field rules.

However, configurations may eventually arrive from JSON, TOML, command-line arguments, environment variables, or HTTP requests. Those boundaries often supply strings or loosely typed dictionaries. In that situation, a validation library such as Pydantic can validate and parse external data into a dataclass-like structure.

Pydantic dataclasses can add constraints such as “greater than zero” and can coerce compatible inputs. That is useful at an external boundary, but coercion is a policy decision: converting "8" to 8 may be convenient for environment variables, while rejecting it may be safer for a strictly controlled internal configuration.

For this course project, the current design is deliberate:

• use standard dataclasses plus explicit semantic validation for experiment code;
• use Pydantic later at API boundaries, where requests are untrusted and structured error responses matter.

Key takeaways

A production configuration is more than a collection of constants. It is a typed, named, validated record of the choices that define a model run.

You have now:

• organized related model, training, and experiment settings as nested dataclasses;
• used frozen=True, slots=True, and kw_only=True to make configuration harder to misuse;
• implemented validation in __post_init__;
• checked both individual value domains and cross-field invariants;
• used immutable tuples for fixed configuration sequences; and
• learned when field(default_factory=...) is necessary for mutable defaults.

Next, you will shift from configuration to data flow: using iterators and generators to stream samples without loading an entire dataset into memory.