"""Typed experiment configuration loaded from YAML.""" from __future__ import annotations from dataclasses import asdict, dataclass from pathlib import Path from typing import Any import yaml MODEL_BACKBONES = ("project", "hf-qwen3") @dataclass(frozen=True) class ModelConfig: """Architecture for the bidirectional token denoiser.""" vocab_size: int mask_token_id: int max_seq_len: int d_model: int n_layers: int n_heads: int d_ff: int dropout: float = 0.1 tie_embeddings: bool = True activation_checkpointing: bool = False use_flex_attention: bool = False compile_backbone: bool = False forbidden_output_token_ids: tuple[int, ...] = (0, 1, 2, 4) backbone: str = "project" pretrained_path: str | None = None # First id of the untrained tail of a pretrained embedding matrix; logits from that id # upward are masked without enumerating hundreds of forbidden ids. forbidden_output_from: int | None = None def __post_init__(self) -> None: if self.vocab_size <= 1: raise ValueError("vocab_size must be greater than one") if self.backbone not in MODEL_BACKBONES: raise ValueError(f"backbone must be one of {MODEL_BACKBONES}") if self.backbone != "project" and self.pretrained_path is None: raise ValueError("pretrained backbones require pretrained_path") if self.forbidden_output_from is not None and not ( 0 < self.forbidden_output_from <= self.vocab_size ): raise ValueError("forbidden_output_from must be inside the vocabulary") if not 0 <= self.mask_token_id < self.vocab_size: raise ValueError("mask_token_id must be inside the vocabulary") if self.max_seq_len <= 0: raise ValueError("max_seq_len must be positive") if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0 or self.d_ff <= 0: raise ValueError("all model dimensions must be positive") if self.d_model % self.n_heads != 0: raise ValueError("d_model must be divisible by n_heads") if not 0.0 <= self.dropout < 1.0: raise ValueError("dropout must be in [0, 1)") if not isinstance(self.activation_checkpointing, bool): raise ValueError("activation_checkpointing must be a boolean") if not isinstance(self.use_flex_attention, bool): raise ValueError("use_flex_attention must be a boolean") if not isinstance(self.compile_backbone, bool): raise ValueError('compile_backbone must be a boolean') forbidden_ids = tuple(self.forbidden_output_token_ids) if any( isinstance(token_id, bool) or not isinstance(token_id, int) for token_id in forbidden_ids ): raise ValueError("forbidden_output_token_ids must contain integers") if len(set(forbidden_ids)) != len(forbidden_ids): raise ValueError("forbidden_output_token_ids must not contain duplicates") if any(not 0 <= token_id < self.vocab_size for token_id in forbidden_ids): raise ValueError("forbidden output token ids must be inside the vocabulary") if self.mask_token_id not in forbidden_ids: raise ValueError("mask_token_id must be a forbidden output token") if self.backbone == "project" and 3 in forbidden_ids: raise ValueError("EOS token id 3 must remain an allowed output token") object.__setattr__(self, "forbidden_output_token_ids", forbidden_ids) @dataclass(frozen=True) class TrainingConfig: """Single-device training settings for the first research iteration.""" train_data: str tokenizer: str output_dir: str = "outputs/run" val_data: str | None = None batch_size: int = 32 gradient_accumulation_steps: int = 1 max_steps: int = 10_000 learning_rate: float = 3e-4 min_learning_rate: float = 3e-5 warmup_steps: int = 500 weight_decay: float = 0.1 grad_clip: float = 1.0 mask_eps: float = 1e-3 seed: int = 1337 device: str = "auto" precision: str = "auto" optimizer: str = "adamw" optimizer_min_8bit_size: int = 4096 optimizer_embedding_32bit: bool = True require_fused_attention: bool = False save_inference_checkpoint: bool = False num_workers: int = 0 log_interval: int = 10 eval_interval: int = 500 eval_batches: int = 20 save_interval: int = 500 keep_last_checkpoints: int = 3 def __post_init__(self) -> None: if self.batch_size <= 0 or self.gradient_accumulation_steps <= 0: raise ValueError("batch sizes must be positive") if self.max_steps <= 0: raise ValueError("max_steps must be positive") if not 0.0 < self.learning_rate: raise ValueError("learning_rate must be positive") if not 0.0 <= self.min_learning_rate <= self.learning_rate: raise ValueError("min_learning_rate must be between zero and learning_rate") if not 0 <= self.warmup_steps < self.max_steps: raise ValueError("warmup_steps must be non-negative and less than max_steps") if self.weight_decay < 0.0: raise ValueError("weight_decay must be non-negative") if self.grad_clip <= 0.0: raise ValueError("grad_clip must be positive") if not 0.0 < self.mask_eps < 1.0: raise ValueError("mask_eps must be in (0, 1)") if self.num_workers < 0: raise ValueError("num_workers must be non-negative") if self.log_interval <= 0 or self.eval_interval <= 0 or self.save_interval <= 0: raise ValueError("log, eval, and save intervals must be positive") if self.eval_batches <= 0: raise ValueError("eval_batches must be positive") if self.keep_last_checkpoints < 0: raise ValueError("keep_last_checkpoints must be non-negative") if self.precision not in {"auto", "float32", "bfloat16", "float16"}: raise ValueError("precision must be auto, float32, bfloat16, or float16") if self.optimizer not in {"adamw", "adamw8bit"}: raise ValueError("optimizer must be adamw or adamw8bit") if self.optimizer_min_8bit_size <= 0: raise ValueError("optimizer_min_8bit_size must be positive") if not isinstance(self.optimizer_embedding_32bit, bool): raise ValueError("optimizer_embedding_32bit must be a boolean") if not isinstance(self.require_fused_attention, bool): raise ValueError("require_fused_attention must be a boolean") if not isinstance(self.save_inference_checkpoint, bool): raise ValueError("save_inference_checkpoint must be a boolean") @dataclass(frozen=True) class ExperimentConfig: model: ModelConfig training: TrainingConfig def to_dict(self) -> dict[str, Any]: return asdict(self) def load_config(path: str | Path) -> ExperimentConfig: """Load and validate an experiment YAML file.""" config_path = Path(path) with config_path.open("r", encoding="utf-8") as handle: raw = yaml.safe_load(handle) if not isinstance(raw, dict) or "model" not in raw or "training" not in raw: raise ValueError("config must contain top-level model and training mappings") return ExperimentConfig( model=ModelConfig(**raw["model"]), training=TrainingConfig(**raw["training"]), )