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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user