| """Run configuration: model shape and training hyperparameters, loaded from YAML.""" |
|
|
| from __future__ import annotations |
|
|
| import dataclasses |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Literal |
|
|
| import yaml |
|
|
|
|
| @dataclass |
| class ModelConfig: |
| d_model: int |
| n_layers: int |
| n_heads: int |
| ffn_hidden: int |
| context_len: int |
| vocab_size: int = 16384 |
| rope_theta: float = 10000.0 |
| norm_eps: float = 1e-5 |
| init_std: float = 0.02 |
|
|
|
|
| @dataclass |
| class TrainConfig: |
| peak_lr: float |
| total_steps: int |
| micro_batch_size: int |
| grad_accum: int |
| precision: Literal["fp32", "fp16"] |
| data_dir: str |
| mix: dict[str, float] |
| max_epochs: dict[str, float] = field(default_factory=dict) |
| warmup_steps: int = 500 |
| min_lr_ratio: float = 0.1 |
| weight_decay: float = 0.1 |
| beta1: float = 0.9 |
| beta2: float = 0.95 |
| grad_clip: float = 1.0 |
| checkpoint_minutes: float = 20.0 |
| eval_every: int = 200 |
| eval_batches: int = 20 |
| seed: int = 1337 |
|
|
|
|
| @dataclass |
| class RunConfig: |
| name: str |
| model: ModelConfig |
| train: TrainConfig |
|
|
|
|
| def _build(cls, data: dict, section: str): |
| known = {f.name for f in dataclasses.fields(cls)} |
| unknown = set(data) - known |
| if unknown: |
| raise ValueError(f"unknown keys in '{section}': {sorted(unknown)}") |
| return cls(**data) |
|
|
|
|
| def load_run_config(path) -> RunConfig: |
| data = yaml.safe_load(Path(path).read_text()) |
| unknown = set(data) - {"name", "model", "train"} |
| if unknown: |
| raise ValueError(f"unknown top-level keys: {sorted(unknown)}") |
| return RunConfig( |
| name=data["name"], |
| model=_build(ModelConfig, data["model"], "model"), |
| train=_build(TrainConfig, data["train"], "train"), |
| ) |
|
|
