code/config.py

"""Single source of truth for the compact 50M dense GPT training run."""

from dataclasses import asdict, dataclass
from pathlib import Path


@dataclass(frozen=True)
class ModelConfig:
    # 32k BPE + tied embeddings. Total model size is ~50.0M parameters.
    vocab_size: int = 32_000
    block_size: int = 512
    n_layers: int = 12
    n_heads: int = 8
    d_model: int = 480
    mlp_hidden: int = 1_352
    dropout: float = 0.0


@dataclass(frozen=True)
class TrainConfig:
    # Optimized profile targeting ~4.5GB VRAM on an RTX 3050 6GB GPU.
    micro_batch_size: int = 4
    gradient_accumulation: int = 8
    learning_rate: float = 3e-4
    weight_decay: float = 0.1
    max_steps: int = 61_000
    warmup_ratio: float = 0.02
    log_every: int = 100
    save_every: int = 1_000
    seed: int = 42


ROOT = Path(__file__).resolve().parent
DATA_DIR = ROOT / ".data"
RAW_DATASET = ROOT / "dataset.jsonl"
CLEAN_DATASET = DATA_DIR / "clean.jsonl"
TOKENIZER_PATH = DATA_DIR / "tokenizer.json"
TOKENS_PATH = DATA_DIR / "tokens.bin"
DATA_METADATA_PATH = DATA_DIR / "metadata.json"
CHECKPOINT_DIR = ROOT / "checkpoints"

MODEL = ModelConfig()
TRAIN = TrainConfig()


def model_config_dict() -> dict:
    return asdict(MODEL)