File size: 1,218 Bytes
7aa9bc4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | from __future__ import annotations
from dataclasses import asdict, dataclass
MODEL_NAME = "Sol Milkshake"
ORGANIZATION = "Sol Labs"
DEPLOYED_PARAMS = 2_990_000
VOCAB_SIZE = 2_048
D_MODEL = 192
MAX_CONTEXT = 2_048
N_PHYSICAL_BLOCKS = 5
N_Q_HEADS = 6
N_KV_HEADS = 2
HEAD_DIM = 32
FFN_HIDDEN = 512
TN_RANK = 26
TIED_EMBEDDING_HEAD = True
XSA_PARAMETER_COUNT = 0
TOKENS_PER_PARAMETER = 544
GLOBAL_TOKENS = 1_626_560_000
GLOBAL_BATCH = 32_768
COMPLETE_UPDATES = 49_638
FINAL_PARTIAL_TOKENS = 22_016
assert DEPLOYED_PARAMS * TOKENS_PER_PARAMETER == GLOBAL_TOKENS
assert COMPLETE_UPDATES * GLOBAL_BATCH + FINAL_PARTIAL_TOKENS == GLOBAL_TOKENS
@dataclass(frozen=True)
class SolConfig:
model_name: str = MODEL_NAME
vocab_size: int = VOCAB_SIZE
d_model: int = D_MODEL
max_context: int = MAX_CONTEXT
n_blocks: int = N_PHYSICAL_BLOCKS
n_q_heads: int = N_Q_HEADS
n_kv_heads: int = N_KV_HEADS
head_dim: int = HEAD_DIM
ffn_hidden: int = FFN_HIDDEN
tn_rank: int = TN_RANK
rope_theta: float = 20_000.0
memory_width: int = 64
memory_slots: int = 32
chunk_size: int = 32
passes: int = 3
dropout: float = 0.0
def as_dict(self) -> dict:
return asdict(self)
|