rzen
234a37f8fd
- src/train.py: AdamW (decay on matrices only), cosine LR w/ warmup, grad clipping, periodic train/val loss, CSV logging, resumable checkpoints. --overfit mode runs the overfit-one-batch sanity check. MPS, fp32. - configs/v1.py: TrainConfig (batch 64, peak LR 6e-4, 20k iters, etc.) - tests/test_train.py: LR schedule + optimizer grouping (3 tests, no run)
51 lines
1.8 KiB
Python
51 lines
1.8 KiB
Python
"""Experiment v1 configuration.
|
|
|
|
One experiment = one config file. Iteration 2 is a copy of this file with
|
|
different numbers, not a code change.
|
|
|
|
v1 is deliberately small (~4.3M parameters): the goal is a complete, legible
|
|
pipeline, not a strong model. It sits below TinyStories' coherence sweet spot
|
|
(~30M params), so some incoherence in the output is expected.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ModelConfig:
|
|
vocab_size: int = 4096 # matches the trained tokenizer
|
|
context_length: int = 256 # tokens of history the model attends over
|
|
n_layers: int = 4
|
|
n_heads: int = 4
|
|
d_model: int = 256 # residual-stream width
|
|
d_ff: int = 1024 # MLP hidden width (4 x d_model)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TrainConfig:
|
|
# Data / batching. tokens per step = batch_size * context_length = 16,384.
|
|
batch_size: int = 64
|
|
|
|
# AdamW optimiser.
|
|
learning_rate: float = 6e-4 # peak LR (reached after warmup)
|
|
min_lr: float = 6e-5 # cosine-decay floor (~10% of peak)
|
|
warmup_iters: int = 200
|
|
max_iters: int = 20_000 # ~327M tokens (~0.6 epoch of the 555M corpus)
|
|
weight_decay: float = 0.1 # applied to matrices only, not norms/biases
|
|
beta1: float = 0.9
|
|
beta2: float = 0.95
|
|
grad_clip: float = 1.0 # clip global grad norm to this
|
|
|
|
# Evaluation, logging, checkpointing (all in optimiser steps).
|
|
eval_interval: int = 500 # how often to measure train/val loss
|
|
eval_iters: int = 100 # batches averaged per loss estimate
|
|
log_interval: int = 20 # how often to print step/loss/lr/speed
|
|
checkpoint_interval: int = 1000
|
|
|
|
seed: int = 1337
|
|
|
|
|
|
# Single source of truth other modules import.
|
|
model = ModelConfig()
|
|
train = TrainConfig()
|