Gahow Wang
15f1e526c7
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>
88 lines
3.2 KiB
Rust
88 lines
3.2 KiB
Rust
//! Tiny-transformer hyperparameters. Host-only (no GPU), always compiled.
|
||
|
||
/// Architecture config for [`crate::TinyTransformer`]. Keep it tiny — T5 is a
|
||
/// correctness bring-up, not a real training run.
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct Config {
|
||
/// Vocabulary size (char-level in the bring-up).
|
||
pub vocab: usize,
|
||
/// Model / residual width. Must equal `n_heads * head_dim`.
|
||
pub dim: usize,
|
||
/// Number of decoder blocks.
|
||
pub n_layers: usize,
|
||
/// Number of attention heads.
|
||
pub n_heads: usize,
|
||
/// Per-head dimension (`dim / n_heads`).
|
||
pub head_dim: usize,
|
||
/// SwiGLU hidden width (gate/up project to this, down projects back).
|
||
pub ffn_hidden: usize,
|
||
/// RMSNorm epsilon.
|
||
pub eps: f32,
|
||
/// RoPE base frequency (theta).
|
||
pub rope_theta: f32,
|
||
}
|
||
|
||
impl Config {
|
||
/// A minimal config used by the bring-up / overfit test.
|
||
pub fn tiny() -> Self {
|
||
let n_heads = 2;
|
||
let head_dim = 16;
|
||
Config {
|
||
vocab: 0, // set by the caller from the char vocab
|
||
dim: n_heads * head_dim,
|
||
n_layers: 2,
|
||
n_heads,
|
||
head_dim,
|
||
ffn_hidden: 64,
|
||
eps: 1e-5,
|
||
rope_theta: 10000.0,
|
||
}
|
||
}
|
||
|
||
/// 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
|
||
+ 2 * self.head_dim // q/k per-head norm gammas
|
||
+ 3 * self.dim * self.dim // q/k/v proj
|
||
+ self.dim * self.dim // out proj
|
||
+ 2 * self.dim * self.ffn_hidden // gate/up proj
|
||
+ self.ffn_hidden * self.dim; // down proj
|
||
self.vocab * self.dim // embedding
|
||
+ self.n_layers * per_layer
|
||
+ self.dim // final norm
|
||
+ self.dim * self.vocab // lm head
|
||
}
|
||
}
|