rzen
9fd56b6063
Controlled study of how the model degrades. Four variants, each = v1 with a single field replaced (dataclasses.replace), so differences are attributable: - small: capacity down (2L/64d, ~0.38M params) - short: training time down (max_iters 20k->2k) - ctx: context window down (256->64) - data: data quantity down (max_train_tokens->1M; new TrainConfig knob + train.py slice) - scripts/compare.py: sample same prompt across all trained configs with a shared seed, print side by side, write reports/compare.md - tests/test_configs.py: enforces one-parameter-at-a-time (only intended fields differ from v1) + small param count (3 tests). Full suite: 29 passing.
52 lines
1.8 KiB
Python
52 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
|
|
max_train_tokens: int | None = None # None = whole train set; set to starve data
|
|
|
|
# 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()
|