diff --git a/crates/xtrain-model/src/config.rs b/crates/xtrain-model/src/config.rs index d5aac59..3eba07a 100644 --- a/crates/xtrain-model/src/config.rs +++ b/crates/xtrain-model/src/config.rs @@ -39,6 +39,38 @@ impl Config { } } + /// Build a config from the architecture knobs, deriving `dim = n_heads * + /// head_dim`. The scaling-run entry (`bin/train`) passes these from CLI so the + /// model size is a tunable ladder rung (v1 = dim256/8L, v2/v3 scale further), + /// instead of a hardcoded tiny config. `eps`/`rope_theta` keep the engine + /// defaults (also what the xserv export reconciles against). + pub fn from_arch( + vocab: usize, + n_heads: usize, + head_dim: usize, + n_layers: usize, + ffn_hidden: usize, + ) -> Self { + Config { + vocab, + dim: n_heads * head_dim, + n_layers, + n_heads, + head_dim, + ffn_hidden, + eps: 1e-5, + rope_theta: 10000.0, + } + } + + /// Transformer-core parameter count: everything except the token embedding and + /// the LM head (the two `vocab × dim` tables). This is the figure the scaling + /// ladder is sized against — the 50257-vocab embed+lm_head adds a fixed ~25M on + /// top that does not reflect model capacity. `num_params() = core + 2·vocab·dim`. + pub fn core_params(&self) -> usize { + self.num_params() - 2 * self.vocab * self.dim + } + /// Total learnable parameter count (for logging / sanity). pub fn num_params(&self) -> usize { let per_layer = 2 * self.dim // 2 rmsnorm gammas