| from __future__ import annotations |
| from dataclasses import asdict, dataclass, fields |
| import json |
| import math |
| from pathlib import Path |
|
|
|
|
| @dataclass(frozen=True) |
| class ModelConfig: |
| vocab_size: int = 32000 |
| layers: int = 12 |
| hidden_size: int = 768 |
| heads: int = 12 |
| intermediate_size: int = 2048 |
| context_length: int = 2048 |
| rope_theta: float = 10000.0 |
| rms_epsilon: float = 1e-5 |
| dropout: float = 0.0 |
| gradient_checkpointing: bool = False |
|
|
| def __post_init__(self): |
| dims=(self.vocab_size, self.layers, self.hidden_size, self.heads, self.intermediate_size, self.context_length) |
| if any(type(value) is not int for value in dims) or min(dims) <= 0: |
| raise ValueError("Model dimensions must be positive") |
| if self.hidden_size % self.heads or (self.hidden_size // self.heads) % 2: |
| raise ValueError("Head dimension must be an even integer") |
| if not all(math.isfinite(value) for value in (self.dropout,self.rms_epsilon,self.rope_theta)) or not 0 <= self.dropout < 1 or self.rms_epsilon <= 0 or self.rope_theta <= 0: |
| raise ValueError("Invalid dropout/RMS epsilon/RoPE theta") |
|
|
| @property |
| def parameter_count(self) -> int: |
| h, f = self.hidden_size, self.intermediate_size |
| return self.vocab_size*h + self.layers*(4*h*h + 3*h*f + 2*h) + h |
|
|
| def to_dict(self): |
| return asdict(self) |
|
|
|
|
| @dataclass(frozen=True) |
| class TrainConfig: |
| batch_size: int = 2 |
| gradient_accumulation: int = 16 |
| max_steps: int = 1000 |
| schedule_total_steps: int | None = None |
| learning_rate: float = 3e-4 |
| min_learning_rate_ratio: float = 0.1 |
| warmup_steps: int = 100 |
| weight_decay: float = 0.1 |
| beta1: float = 0.9 |
| beta2: float = 0.95 |
| grad_clip: float = 1.0 |
| seed: int = 1337 |
| precision: str = "auto" |
| checkpoint_every: int = 100 |
| eval_every: int = 100 |
| eval_batches: int = 10 |
| deterministic: bool = False |
| max_full_checkpoints_per_session: int = 3 |
| checkpoint_interval_minutes: float = 0.0 |
| checkpoint_token_milestones: tuple[int, ...] = () |
|
|
| def __post_init__(self): |
| for name in ("batch_size", "gradient_accumulation", "max_steps", "checkpoint_every", "eval_every", "eval_batches"): |
| if type(getattr(self,name)) is not int or getattr(self, name) <= 0: |
| raise ValueError(f"{name} must be positive") |
| if self.schedule_total_steps is not None and (type(self.schedule_total_steps) is not int or self.schedule_total_steps < self.max_steps): |
| raise ValueError("Explicit LR schedule horizon must cover this phase's max_steps") |
| if not isinstance(self.checkpoint_token_milestones,(list,tuple)): |
| raise ValueError("Checkpoint token milestones must be a list/tuple") |
| milestones=tuple(self.checkpoint_token_milestones) |
| if any(type(v) is not int or v<=0 for v in milestones) or tuple(sorted(set(milestones))) != milestones: |
| raise ValueError("Checkpoint milestones must be strictly increasing positive integers") |
| if len(milestones)>self.max_full_checkpoints_per_session-2: |
| raise ValueError("Milestones must leave start and final checkpoint slots") |
| object.__setattr__(self,"checkpoint_token_milestones",milestones) |
| if type(self.seed) is not int or self.seed < 0: |
| raise ValueError("Seed must be a nonnegative integer") |
| if type(self.warmup_steps) is not int or not 0 <= self.warmup_steps <= self.max_steps: |
| raise ValueError("warmup_steps must be between zero and max_steps") |
| if not all(math.isfinite(value) for value in (self.learning_rate,self.grad_clip,self.weight_decay,self.min_learning_rate_ratio,self.beta1,self.beta2)): |
| raise ValueError("Optimizer values must be finite") |
| if self.learning_rate <= 0 or self.grad_clip <= 0 or self.weight_decay < 0: |
| raise ValueError("Invalid optimizer parameters") |
| if not 0 <= self.min_learning_rate_ratio <= 1 or not 0 <= self.beta1 < 1 or not 0 <= self.beta2 < 1: |
| raise ValueError("Invalid optimizer ratios") |
| if type(self.max_full_checkpoints_per_session) is not int or not 2 <= self.max_full_checkpoints_per_session <= 3: |
| raise ValueError("Keep 2–3 full checkpoint slots per session") |
| if not math.isfinite(self.checkpoint_interval_minutes) or self.checkpoint_interval_minutes < 0: |
| raise ValueError("Checkpoint interval minutes must be finite and nonnegative (zero means adaptive)") |
| if self.precision not in {"auto", "float32", "bfloat16", "float16"}: |
| raise ValueError("Unsupported precision") |
|
|
| def to_dict(self): |
| return asdict(self) |
|
|
|
|
| def load_config(path: str | Path): |
| with open(path) as f: |
| obj = json.load(f) |
| unknown = set(obj) - {"model", "training", "description", "provisional"} |
| if unknown: |
| raise ValueError(f"Unknown config keys: {sorted(unknown)}") |
| return ModelConfig(**obj["model"]), TrainConfig(**obj["training"]) |
|
|
