| from dataclasses import dataclass | |
| class ARConfig: | |
| """All the numbers that define the AR model's shape. | |
| Swap these to build a different-sized model with the same code | |
| (e.g. a tiny version for debugging). | |
| """ | |
| vocab_size: int = 151936 # size of the tokenizer vocabulary | |
| dim: int = 2560 # hidden size (width of the residual stream) | |
| n_layers: int = 36 # number of decoder blocks stacked | |
| n_heads: int = 32 # query heads | |
| n_kv_heads: int = 8 # key/value heads (GQA: 32 queries share 8 kv) | |
| head_dim: int = 128 # size of each attention head | |
| hidden_dim: int = 9728 # inner width of the SwiGLU MLP | |
| max_seq_len: int = 40960 # longest sequence the model supports | |
| rope_theta: float = 1_000_000.0 # RoPE base frequency | |
| rms_norm_eps: float = 1e-6 # RMSNorm epsilon | |
| tie_embeddings: bool = True # share input embedding and output head weights | |
| # a very small variant for fast offline code testing (tiny numbers) | |
| def tiny_config() -> ARConfig: | |
| return ARConfig( | |
| vocab_size=1000, dim=128, n_layers=2, | |
| n_heads=4, n_kv_heads=2, head_dim=32, | |
| hidden_dim=512, max_seq_len=256, | |
| ) | |
| # a ~0.5B variant for validating the training loop on a single GPU | |
| # (same architecture shape as the full model, just narrower and shallower) | |
| def small_config() -> ARConfig: | |
| return ARConfig( | |
| vocab_size=151936, dim=1024, n_layers=24, | |
| n_heads=16, n_kv_heads=4, head_dim=128, # keeps the 4:1 GQA ratio | |
| hidden_dim=2816, max_seq_len=40960, | |
| ) |