rzen
e708a5b886
- src/model.py: GPT written module by module — RMSNorm, explicit causal self-attention (scores/mask/softmax/weighted-sum), GELU MLP, pre-norm residual blocks, learned positions, weight-tied head. GPT-2 style init. - configs/v1.py: frozen ModelConfig dataclass (4L/256d/4h/256ctx/4096vocab) - tests/test_model.py: exact param count (4,262,144), forward shapes, init loss ~= ln(vocab), weight tying, causal-masking check (5 tests, forward-only) - Verified forward pass runs on MPS.
26 lines
829 B
Python
26 lines
829 B
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)
|
|
|
|
|
|
# The single source of truth other modules import.
|
|
model = ModelConfig()
|