train: parameterize model size (scaling ladder)

Add Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn) so the model
size is a tunable rung instead of a hardcoded tiny config, and Config::core_params()
(num_params minus the two vocab×dim tables) — the figure the ladder is sized
against (the 50257-vocab embed+lm_head adds a fixed ~25M that is not capacity).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-15 18:34:39 +08:00
parent 8981cf7982
commit 15f1e526c7

View File

@@ -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